Source-linked AI summary
Generating Long Sequences with Sparse Transformers
Rewon Child, Scott Gray, Alec Radford, Ilya Sutskever
TL;DR
Dense self-attention is powerful for sequence modeling but becomes impractical on long sequences because its memory and computation grow quadratically. The paper introduces Sparse Transformers with sparse attention factorizations and supporting architectural changes, achieving state-of-the-art density modeling and generation across natural language, raw audio, and natural images.
Problem
Autoregressive models must learn complex long-range dependencies, while standard self-attention has quadratic memory and computational requirements that exclude long sequences.
Method
Sparse Transformers factorize attention into sparse operations and add changes for deep-network training, efficient kernels, memory-saving recomputation, and diverse data representations.
Results
State-of-the-art compression and generation are achieved for natural language, raw audio, and natural images.
Takeaways & Limitations
Sparse Transformers provide a single architecture for density modeling of long sequences across images, text, and audio, with globally coherent samples.
Takeaways & Limitations
For data without periodic structure such as text, strided attention can fail to route information properly because spatial coordinates may not match future relevance.
Abstract
from arXiv · showhide
Transformers are powerful sequence models, but require time and memory that grows quadratically with the sequence length. In this paper we introduce sparse factorizations of the attention matrix which reduce this to $O(n \sqrt{n})$. We also introduce a) a variation on architecture and initialization to train deeper networks, b) the recomputation of attention matrices to save memory, and c) fast attention kernels for training. We call networks with these changes Sparse Transformers, and show they can model sequences tens of thousands of timesteps long using hundreds of layers. We use the same architecture to model images, audio, and text from raw bytes, setting a new state of the art for density modeling of Enwik8, CIFAR-10, and ImageNet-64. We generate unconditional samples that demonstrate global coherence and great diversity, and show it is possible in principle to use self-attention to model sequences of length one million or more.
1. Introduction
The paper targets long-range sequence modeling, where autoregressive conditionals are difficult to learn and dense self-attention becomes impractical because its requirements grow quadratically with sequence length. It introduces Sparse Transformers, combining sparse attention factorizations with architectural and computational changes for efficient modeling across data types.
- Motivation: Autoregressive models must learn complex long-range dependencies in conditional distributions over sequences.These dependencies make conditional-distribution modeling challenging despite the success of neural autoregressive models.
- Applications: The same self-attention-based architecture is applied to audio, images, and text, including unconditional samples from ImageNet 64 and classical music.The figure reports sample lengths of 12,288 and 65,536 at softmax temperature 1.0.
- Motivation: Quadratic memory and computational requirements exclude standard self-attention networks from long-sequence modeling.Each self-attention layer computes n weightings for each of n sequence elements.
- Sparse Transformers: Sparse attention factorizations scale as O(n √n) with sequence length without sacrificing performance.The factorizations separate full attention into several faster attention operations that can approximate dense attention.
- Sparse Transformers: Sparse Transformers add a restructured residual block and initialization, sparse attention kernels, and attention-weight recomputation.These changes target deeper-network training, efficient subset computation, and reduced memory usage.
- Results: The augmented models achieve state-of-the-art compression and generation for natural language, raw audio, and natural images.The authors describe the architecture as simple and potentially useful for many problems of interest.
2. Related Work
The related work develops alternative ways to scale autoregressive generation and attention across images, text, and audio. Sparse Transformers are presented as a simpler, cross-domain approach that can also combine with several existing techniques.
- Generative modeling: Prior image-generation methods use conditional pixel independence, ordering with multiscale upsampling, or local attention blocks.These approaches target parallel generation, high-fidelity samples, or applying Transformers to images.
- Generative modeling: Prior text-generation work introduces state reuse through a memory mechanism for modeling long-term dependencies.The cited approach is Transformer-XL.
- Generative modeling: Prior audio work includes autoregressive neural generation methods related to the paper’s cross-domain modeling setting.The related-work discussion places audio alongside image and text scaling techniques.
- Positioning: Sparse Transformers are simpler than many related techniques and apply equally across images, text, and audio.The paper also states that many existing techniques are orthogonal and could be used in conjunction with its approach.
- Attention efficiency: Other attention-efficiency research explores chunking, fixed-length representations, or multiple attention hops.These methods address efficiency or repeated attention processing outside generative modeling.
3. Background
The paper formulates autoregressive generation as predicting each sequence token from its preceding tokens. Images, text, and audio are represented as discrete-token sequences, and a decoder-only Transformer is trained by maximizing data log-probability.
- Autoregressive modeling: Autoregressive generation models a sequence joint probability as a product of conditional distributions parameterized by a network θ.Each conditional predicts the next element given the preceding sequence elements.
- Token representation: Images, text, and audio are treated as sequences of discrete tokens, typically raw bytes.The network outputs a categorical distribution over the vocabulary for the next token.
- Training: The training objective is to maximize the log-probability of the data with respect to θ.This objective trains the model to assign high probability to observed token sequences.
- Transformer baseline: A decoder-only Transformer uses multihead self-attention over the sequence followed by dense transformations at each sequence element.The self-attention computation becomes intractable as sequence length grows because it computes n weightings for each of n elements.
- Sparse Transformers: The paper modifies the Transformer architecture to make it more suitable for modeling long sequences.
4. Factorized Self-Attention
Sparse Transformers factorize self-attention across several steps to reduce computation while preserving global signal propagation. The paper evaluates strided and fixed patterns, balancing efficiency, connectivity, and expressivity across domains.
- Motivation: Learned attention in full Transformers is often sparse, but also includes local, global, and data-dependent patterns that fixed sparsity may not reproduce.On CIFAR-10, early layers resembled convolution, some layers factorized row and column attention, and later layers showed highly selective access.
- Factorized self-attention: Factorized attention separates full self-attention into sequential steps, allowing global connectivity across positions while reducing effective computation to O(n p√n).The approach restricts each head to subsets of indices while requiring every input position to connect to every future output within at most p + 1 steps.
- Evaluation: The factorized patterns preserve global context and improve efficiency, but cannot learn exactly the same mappings as the richer attention patterns learned by standard Transformers.The paper therefore empirically validates these patterns across tasks rather than assuming equivalence to full attention.
- Two-dimensional factorized attention: For two-dimensional data, strided attention combines local attention over previous l locations with attention to every lth location, with l chosen near √n.This pattern is called strided attention and can align naturally with images or some types of music.
- Two-dimensional factorized attention: Strided attention can fail on text when spatial coordinates do not correlate with future relevance, so fixed attention uses summary cells to propagate information forward.Specific cells summarize previous locations and make those summaries available to later cells.
- Two-dimensional factorized attention: Fixed attention with c = 1 limits expressivity, while choosing c ∈ {8, 16, 32} performs well for typical l ∈ {128, 256} but increases computation by c relative to strided attention.Multiple heads perform better when attending to distinct subblocks within each block.
5. Sparse Transformer
The Sparse Transformer modifies Transformer attention and residual blocks to support factorized attention, deeper networks, and memory-efficient training. It applies these changes to diverse data types while preserving long-range modeling capabilities.
- Factorized attention heads: Sparse Transformer attention integrates factorized patterns across residual blocks, including separate, interleaved, merged, and multi-head arrangements.The current residual-block index and number of factorized heads determine which attention pattern is used.
- Factorized attention heads: Multiple attention heads can work well, while sequential computation is more worthwhile when attention dominates runtime for extremely long sequences.Multi-head computation keeps parameter counts invariant by reducing the dimensions of internal weight matrices per head.
- Scaling to hundreds of layers: The architecture uses pre-activation residual blocks with Layer Normalization and positionwise feedforward networks, plus initialization scaled by 1/√(2N) for deeper networks.The scaling keeps the ratio of input embedding scale to residual-block scale invariant across network depths.
- Modeling diverse data types: Learned embeddings encode either data structure or factorized attention patterns, with data embeddings for image locations and two-dimensional attention embeddings for text and audio.The paper reports that these learned embeddings were important for model performance.
- Saving memory by recomputing attention weights: Recomputing attention and feedforward activations during backpropagation enables dense-attention networks with hundreds of layers on sequences of length 16,384.The approach trades additional computation during gradient calculation for substantially lower memory use.
6. Training
Training uses stability-oriented optimization and precision choices, with fixed embedding and transformation dimensions that control model capacity and efficiency.
- Optimization and hardware: The models use Adam with 5,000-step linear warmup, gradient clipping at 1.0, weight decay of 0.01, and cosine learning-rate decay.Training was conducted on 8 V100 GPUs unless otherwise noted.
- Model dimensions: Embeddings have constant dimension d, while feed-forward transformations usually expand inputs to 4d or 2d in half-size configurations.Default linear transformations retain the embedding dimension, except within the feed-forward network.
- Initialization: Weights and biases use specified Gaussian or zero initialization schemes, including zero-initialized output-logit weights and fan-in-scaled internal weights.Position embeddings use a separate Gaussian scale involving d and the number of embeddings.
7. Experiments
Experiments show Sparse Transformers achieve strong density modeling across images, text, and audio while exploiting sparse patterns for efficiency and long-context modeling. Results include state-of-the-art compression, improved losses over dense attention in several comparisons, coherent long-context generation, and scaling beyond one million timesteps.
- Efficiency: Sparse patterns ran faster and achieved lower loss than full attention on datasets where both could be compared.The authors suggest either a useful inductive bias from sparsity or an optimization issue with full attention.
- CIFAR-10: 2.80 bits per dim on CIFAR-10 surpassed the previous 2.85 state of the art.The strided pattern also reached 2.82 bits per dim faster than dense attention.
- Text: 0.99 bits per dim on EnWik8 surpassed the similarly sized Transformer-XL result of 1.03 and matched a model with more than twice the parameters.Fixed sparse patterns surpassed dense attention, while strided attention performed poorly.
- Text: EnWik8 performance increased monotonically with longer contexts, reaching 12,160 of the 12,288 training tokens.This suggests the model incorporates long-term dependencies.
- ImageNet 64x64: 3.44 bits per dim on downsampled ImageNet improved on the previous 3.52 result.The model used a 48-layer strided Sparse Transformer with 152 million parameters.
- ImageNet 64x64: Unconditional ImageNet samples showed no visible sparsity artifacts and exhibited long-term structure in most images.Samples were generated at softmax temperature 1.0, including from a 300-million-parameter model.
- Long-context audio: Factorized self-attention scaled to sequences over 1 million timesteps, although the model had only 3 million parameters.For classical audio, 65,536-step samples showed global coherence and varied playing styles, while quality degraded at greater lengths.
8. Conclusion
Sparse Transformers achieve equivalent or better long-sequence density modeling than standard Transformers with significantly fewer operations. The approach reaches state-of-the-art performance in images and text, adapts to raw audio, and produces globally coherent samples.
- Conclusion: Sparse Transformers matched or exceeded standard Transformers on long-sequence density modeling while requiring significantly fewer operations.The reported performance was state of the art for images and text and adaptable to raw audio.
- Conclusion: The models demonstrated long-term context use and generated globally coherent samples.