Source-linked AI summary
Fast Transformers with Clustered Attention
Apoorv Vyas, Angelos Katharopoulos, François Fleuret
TL;DR
Transformers’ self-attention is powerful for sequence modeling but scales quadratically with sequence length, limiting large-sequence applications. This paper introduces clustered attention, which groups similar queries and computes attention at cluster centroids, with refinements for high-attention keys. The approach has linear complexity for a fixed cluster count and improves equal-budget speech-recognition performance while approximating pretrained models with 25 clusters and little or no performance loss.
Problem
Self-attention has quadratic complexity with sequence length, limiting transformers’ applicability to long sequences.
Method
Clustered attention groups similar queries using clustering, computes attention once per cluster, and separately considers keys with the highest attention.
Results
Clustered attention consistently outperforms vanilla attention at equal computational budgets and approximates pretrained transformers with 25 clusters without performance loss on GLUE and nearly so on SQuAD.
Takeaways & Limitations
The method provides linear asymptotic complexity while retaining close performance to full attention and enabling efficient approximation of pretrained models.
Takeaways & Limitations
Self-attention’s underlying limitations remain for approaches that improve effective sequence length without changing its complexity.
Abstract
from arXiv · showhide
Transformers have been proven a successful model for a variety of tasks in sequence modeling. However, computing the attention matrix, which is their key component, has quadratic complexity with respect to the sequence length, thus making them prohibitively expensive for large sequences. To address this, we propose clustered attention, which instead of computing the attention for every query, groups queries into clusters and computes attention just for the centroids. To further improve this approximation, we use the computed clusters to identify the keys with the highest attention per query and compute the exact key/query dot products. This results in a model with linear complexity with respect to the sequence length for a fixed number of clusters. We evaluate our approach on two automatic speech recognition datasets and show that our model consistently outperforms vanilla transformers for a given computational budget. Finally, we demonstrate that our model can approximate arbitrarily complex attention distributions with a minimal number of clusters by approximating a pretrained BERT model on GLUE and SQuAD benchmarks with only 25 clusters and no loss in performance.
1 Introduction
Transformers use self-attention to model long sequences effectively, but its quadratic cost limits their use on long inputs. Clustered attention approximates self-attention by grouping similar queries, achieving linear complexity for a fixed number of clusters and strong benchmark performance.
- Self-attention helps transformers handle long sequences without the vanishing-gradient problem inherent in RNNs.
- Quadratic self-attention complexity limits transformers’ applicability to long sequences because each query attends to all keys.
- Prior approaches either reduce self-attention’s asymptotic complexity by limiting keys or support longer sequences without changing that complexity.
- Clustered attention groups similar queries using locality-sensitive hashing and K-Means, computing attention once per cluster for linear complexity with a fixed number of clusters.
- Clustered attention consistently outperforms vanilla attention at equal computational budgets on two speech-recognition datasets and approximates pretrained BERT with 25 clusters without performance loss.
2 Related Work
Related work addresses long-sequence transformers by reducing attention computation, restricting context, or extending effective sequence length. Clustered attention instead trades a small approximation error for linear asymptotic complexity while automatically grouping similar queries.
- Long-sequence transformer methods either improve attention’s asymptotic complexity or make longer sequences possible without changing attention’s quadratic complexity.
- Earlier methods aggregate inputs or restrict accessible attention elements, but fixed aggregation or monotonic access can limit expressivity or flexibility.
- Adaptive Attention Span and Transformer-XL reduce requirements or increase effective sequence length, but do not remove the underlying self-attention limitations.
- Clustered attention trades a small full-attention approximation error for improved linear asymptotic complexity, making long-sequence processing possible.
- Fixed local and strided attention factorization can require two layers for global access and is data-independent, whereas clustered attention groups similar queries automatically.
- Set Transformers achieve linear complexity with fixed inducing points, while Reformer uses LSH-based groups but constrains queries and keys to be equal.
3 Scaling Attention with Fast Clustering
Clustered attention reduces self-attention cost by grouping similar queries and computing attention for cluster centroids, while top-k key refinement improves the approximation. For a fixed number of clusters, the resulting computation is linear in sequence length.
- Clustered Attention: Queries are partitioned into C clusters, and attention is computed once per cluster centroid instead of separately for every query.The resulting clustered attention and value aggregation are broadcast to queries belonging to the same cluster.
- Approximation Quality: For nearby queries, the attention approximation error can be bounded by their Euclidean distance from the cluster centroid.This follows from the similarity of attention distributions for queries that are close in Euclidean space.
- Fast Clustering: LSH followed by K-Means in Hamming space accelerates the construction of query clusters.The proposed clustering procedure uses sign random projections and has complexity O(NCL + CBL + NDkB).
- Improved Approximation: Top-k refinement recomputes exact query-key dot products for the highest-attention keys within each cluster.The refined attention redistributes the corresponding probability mass using these exact dot products while retaining clustered attention for other keys.
- Improved Approximation: The top-k correction adds O(Nk max(Dk, Dv)) to the clustered-attention complexity.The added work consists of two sparse dot products per query involving the selected keys and values.
- Approximation Quality: Improved clustered attention always approximates full attention better than basic clustered attention in L1 distance.The paper states this conclusion from Proposition 2 and its associated equation.
4 Experiments
Experiments evaluate clustered attention on speech recognition and pretrained-model approximation, comparing computational cost, accuracy, convergence, and attention-approximation quality across transformer variants.
- Speed–accuracy trade-off: Improved clustered attention outperforms baselines at equal computational budgets on the WSJ and Switchboard speech-recognition tasks.The experiments compare full, Reformer, clustered, and improved clustered variants under matched computation.
- Approximation quality: Improved clustered attention achieves the lowest phone error rate in every WSJ approximation comparison.The comparison evaluates models trained with one attention implementation and inferred with another.
- Approximation quality: Improved clustered attention captures long-tail attention better than retaining only each query’s 32 highest-attention keys.The oracle-top variant has significantly higher phone error rate in all reported comparisons.
- Convergence behaviour: Clustered attention is more than two times faster per epoch than full attention, while improved clustered attention is faster both per epoch and in total convergence time.The convergence comparison reports test PER, epoch time, and total wall-clock convergence time.
- Switchboard: On Switchboard, improved clustered attention is more than 2 percentage points better in WER than full attention at an approximately 50-second forward-pass budget.Both clustered variants are also faster per epoch and in total wall-clock training time on the longer-sequence Switchboard data.
- RoBERTa approximation: With 25 clusters, improved clustered attention matches full performance on all evaluated GLUE tasks and is only marginally worse on SQuAD.The 25-cluster setting is less than 20% of GLUE sequence length and 10% of SQuAD sequence length; full attention is faster for these short sequences.
5 Conclusions
The paper presents clustered attention as a lower-cost approximation to vanilla transformer attention that preserves linear asymptotic complexity and can approximate pretrained full-attention models.
- Conclusion: Clustered attention approximates vanilla transformers with significantly lower computational requirements while retaining linear asymptotic complexity.The method is reported as up to 2× faster during training and inference with minimal performance loss.
- Conclusion: The method efficiently approximates pretrained full-attention models while retaining linear asymptotic complexity.This capability is presented as a contrast with recent fast transformer variations.
- Future directions: The authors identify long-text masked language modeling as an important direction for contexts longer than the commonly used 512 tokens.They also mention applications including music generation and scene flow estimation.
Broader Impact
Clustered attention reduces transformer computational requirements and can reduce GPU training time, energy consumption, and associated CO2 emissions for long sequences.
- Broader impact: For long sequences, clustered attention can reduce GPU training time by almost 50%, with corresponding reductions in energy consumption and CO2 emissions.The paper connects the reduced computational requirements to use on embedded and resource-constrained devices.
Supplementary Material for Fast Transformers with Clustered Attention
Clustered attention groups similar queries, computes attention for cluster centroids, and broadcasts the resulting values to individual queries. Improved clustered attention refines this approximation by recomputing attention on each cluster’s top-k keys using sparse dot products.
- Clustered attention: K-means clusters queries, centroid attention is computed once, and centroid values are broadcast to each query.The supplementary example uses 8 queries and 3 clusters.
- Complexity: Clustered attention requires O(NCD_v) operations for centroid values, compared with O(NkD_v) for per-query top-k values.Here N is sequence length, C is the number of clusters, k is the number of selected keys, and D_v is value dimensionality.
- Improved clustered attention: Top-k key selection improves the approximation by computing exact sparse query–key dot products for selected keys.The top-k keys are stored in T, which marks selected keys for each cluster.
- Improved clustered attention: The improved output combines a top-k contribution with a clustered-attention contribution from the remaining keys.The top-k term uses corresponding values, while the remaining-key term uses clustered attention weights.
- Approximation quality: The improved attention matrix differs from clustered attention only on the selected top-k keys, where its approximation error is analyzed.The proposition focuses the error comparison on those keys.
C.1 Time and Memory Benchmark
The benchmark measures per-element GPU time and memory during forward/backward passes across increasing sequence lengths. Clustered variants scale linearly, use less per-sample memory, and become faster than vanilla attention at sufficiently long sequences.
- Scaling: All non-full models scale linearly with sequence length, while vanilla transformer scales quadratically and exceeds GPU memory beyond 2^13 elements.The benchmark evaluates per-element time and memory over sequence lengths from N = 2^9 to N = 2^15.
- Runtime and memory: Clustered attention becomes faster than vanilla transformer at 1000 elements, while improved clustered attention surpasses it at 2000 elements.Both clustered variants also achieve better per-sample memory usage than the compared methods.
- Runtime and memory: Per-sample memory is lower for clustered and improved clustered attention because they avoid storing intermediate results for multiple hashing rounds.The comparison specifically contrasts them with Reformer’s hashing-round storage requirements.
C.2 Ablation on clusters and sequence length
The synthetic masked-copy task tests how performance changes with sequence length, cluster count, and hashing rounds. Improved clustered attention achieves perfect accuracy across every tested cluster-count and sequence-length combination.
- Task: The masked-copy task evaluates whether models can reconstruct repeated sequences after randomly masking tokens.Solving it requires attending to corresponding tokens in two identical sequence halves.
- Results: Improved clustered attention achieves perfect accuracy for every tested combination of cluster count and sequence length.The tested clustered models use 15, 30, 60, and 100 clusters across multiple sequence lengths.
- Results: Standard clustered attention and Reformer require more clusters or hashing rounds as sequence length increases.Vanilla transformer solves the task perfectly at all tested sequence lengths.
C.3.1 Wall Street Journal
The Wall Street Journal experiment compares full, clustered, improved clustered, and Reformer attention variants during training. Improved clustered attention converges fastest in one setup and achieves comparable wall-clock convergence to full attention while clustered variants outperform Reformer variants.
- Convergence: Improved clustered attention converges faster than clustered attention and is the only variant with comparable wall-clock convergence to full attention.The comparison accounts for the differing batch sizes used by full and clustered models.
- Convergence: Clustered transformers significantly outperform Reformer variants in wall-clock convergence.The reported comparison concerns the training-loss convergence curves.
- Convergence: In the second setup, improved clustered attention has the fastest convergence, while clustered attention still requires less overall training time than full attention.The latter is attributed to earlier overfitting on validation WER by the full model.
C.4 RoBERTa Approximation
With only 25 clusters, improved clustered attention closely matches full attention on a complex SQuAD question, while clustered attention produces qualitatively different patterns.
- 25 clusters let improved clustered attention approximate full attention closely, even for complicated and sparse attention distributions.The comparison uses a randomly selected SQuAD question-context tuple.
- Improved clustered and full attention identify the same highest-weight tokens for almost all question tokens.This similarity further supports the approximation of complex attention patterns.
- Clustered attention fails to capture the displayed complex attention distribution, unlike improved clustered attention.The figure comparison describes clustered attention as qualitatively different from full attention.