Source-linked AI summary
Fast Transformer Decoding: One Write-Head is All You Need
Noam Shazeer
TL;DR
Incremental Transformer inference is often slow because autoregressive dependencies prevent parallelization and repeatedly loading large keys and values consumes memory bandwidth. The paper introduces multi-query attention, sharing keys and values across heads to reduce this cost. Experiments show much faster decoding with only minor quality degradation, including a highest BLEU score of 28.5 with beam-4 decoding.
Problem
Incremental Transformer inference is slowed by unavoidable sequential generation and the memory-bandwidth cost of repeatedly loading large keys and values tensors.
Method
Multi-query attention retains multiple attention heads but shares one set of keys and values across them, reducing incremental memory requirements.
Results
28.5 BLEU was achieved by the multi-query model with beam-4 decoding, while performing similarly to the baseline.
Takeaways & Limitations
Multi-query attention offers lower memory-bandwidth requirements for incremental inference and may enable wider adoption of attention-based sequence models in inference-critical applications.
Takeaways & Limitations
The speed evaluation used padding and masking for fixed shapes, so incrementally growing tensors could save time near the beginning of sequences.
Abstract
from arXiv · showhide
Multi-head attention layers, as used in the Transformer neural sequence model, are a powerful alternative to RNNs for moving information across and between sequences. While training these layers is generally fast and simple, due to parallelizability across the length of the sequence, incremental inference (where such paralleization is impossible) is often slow, due to the memory-bandwidth cost of repeatedly loading the large "keys" and "values" tensors. We propose a variant called multi-query attention, where the keys and values are shared across all of the different attention "heads", greatly reducing the size of these tensors and hence the memory bandwidth requirements of incremental decoding. We verify experimentally that the resulting models can indeed be much faster to decode, and incur only minor quality degradation from the baseline.
1 Introduction
Transformer offers a popular alternative to recurrent sequence models, but incremental inference is slowed by repeatedly reloading large attention keys and values. The paper proposes multi-query attention to improve inference speed with minor quality degradation.
- Incremental Transformer inference is limited by the memory bandwidth required to reload large keys and values tensors.
- Multi-query attention shares keys and values across attention heads to reduce tensor size and improve inference speed.
- The proposed variation is evaluated for real performance gains while maintaining high model quality.
2 Background: Neural Attention
Neural attention computes outputs by weighting value vectors according to query-key comparisons, while Transformer multi-head attention runs multiple projected heads in parallel. During autoregressive incremental inference, data dependencies prevent parallel queries, making repeated key-value memory access a bottleneck.
- 2 Background: Neural Attention: Neural attention produces an output by weighting value vectors according to comparisons between one query and multiple keys.
- 2.2 Multi-head Attention: Transformer multi-head attention runs h parallel heads with separate learned projections for queries, keys, values, and outputs.
- 2.3.1 Performance Analysis of Batched Multi-head Attention: Batched multi-head attention performs Θ(bnd2) arithmetic operations and accesses O(bhn2 + bnd + d2) memory under simplifying assumptions.
- 2.3.1 Performance Analysis of Batched Multi-head Attention: The memory-access-to-computation ratio is O(1/bn), which suits hardware whose computational capacity can greatly exceed memory bandwidth.
- 2.4 Multihead Attention (Incremental): Autoregressive generation prevents parallel computation because each generated token affects the next self-attention query.
- 2.4.1 Performance Analysis: Across n incremental calls, key-value memory access is Θ(bn2d + nd2), making memory bandwidth a bottleneck when n ≈ d or b ≈ 1.
- 2.4.1 Performance Analysis: Reducing sequence length or attended positions can shrink key-value tensors, while the paper proposes removing their heads dimension as an orthogonal approach.
3 Multi-Query Attention
Multi-query attention preserves separate query and output projections but shares one set of keys and values across heads. This reduces incremental memory access while retaining high quality and producing experimentally verified performance gains.
- 3 Multi-Query Attention: Multi-query attention differs from multi-head attention by sharing a single set of keys and values across all heads.
- 3 Multi-Query Attention: In batched multi-query attention, K and V omit the heads dimension while Q and output projections retain head-specific structure.
- 3 Multi-Query Attention: Incremental multi-query attention stores shared K and V tensors with shapes [b, m, k] and [b, m, v], rather than per-head tensors.
- 3.1 Performance Analysis for Incremental Multi-Query Attention: Across n calls, multi-query attention has Θ(bnd2) arithmetic operations and Θ(bnd + bn2k + nd2) memory access.
- 3.1 Performance Analysis for Incremental Multi-Query Attention: The reduced key-value term is expected to improve incremental generation substantially at large batch sizes, while model quality remains high.
4 Experiments and Results
Experiments compare parameter-matched multi-query Transformers with baseline and alternative reduced-dimension models on translation and language modeling. Multi-query maintains near-baseline quality while substantially reducing incremental decoding cost.
- Experimental setup: The experiments cover WMT 2014 English-German translation and the Billion-Word language-modeling benchmark with parameter-matched model variations.The translation baseline and variations have 211 million parameters, while the language-modeling models have 192 million.
- Experimental setup: Multi-query replaces all attention layers with shared-key/value attention and widens feed-forward layers to match the baseline parameter count.The encoder self-attention, decoder self-attention, and encoder-decoder attention layers are replaced.
- Model quality: The Billion-Word results show the same pattern: multi-query is slightly worse than baseline but significantly better than alternatives reducing h, d_k, or d_v.Models were evaluated by per-word perplexity on the development set.
- Model quality: Multi-query is slightly worse than the translation baseline on BLEU and perplexity but closer to it than alternatives that reduce h, d_k, or d_v.The comparison uses dev-set BLEU and per-subword-token perplexity.
- Model quality: 28.5 BLEU was the highest test-set score for multi-query with beam-4 decoding, while its overall performance remained similar to the baseline.Test evaluation used greedy decoding and beam search with beam 4 and α = 0.6.
- Speed: Training time changed from 13.2µs per input-token-plus-target-token for the baseline to 425ms per training step for multi-query, compared on 32,768-token steps.The corresponding baseline and multi-query step times were 433ms and 425ms.
- Speed: Decoder inference fell from 46µs to 3.8µs per token with multi-query, while encoder inference changed from 1.7µs to 1.5µs per token.These are amortized costs for incremental greedy inference on sequences of length 128.
- Speed: Fixed-shape padding and masking made every decoding step use the maximum sequence or window length, potentially overstating costs near sequence beginnings.Growing tensors incrementally could save time near the beginning of a sequence.
5 Conclusion
The paper proposes multi-query attention as an alternative to multi-head attention with lower memory-bandwidth requirements during incremental inference. The authors argue that this can support wider adoption of attention-based sequence models in inference-performance-critical applications.
- Conclusion: Multi-query attention is presented as an alternative to multi-head attention with much lower memory-bandwidth requirements in the incremental setting.The stated application scope is inference-performance-critical attention-based sequence models.