Fast Transformer Decoding: One Write-Head is All You Need

By Noam Shazeer. 2019

825 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Model Architectures, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Model Architectures

In Efficiency & Compression

Citations in the corpus

PaLM: Scaling Language Modeling with Pathways (2022, background): “We have found that this has a neutral effect on model quality and training speed (Shazeer, 2019), but results in a significant cost savings at autoregressive decoding time.”
Llama 2: Open Foundation and Fine-Tuned Chat Models (2023, methodology): “Either the original multi-query format with a single KV projection (MQA, Shazeer, 2019) or a grouped-query attention variant with 8 KV projections (GQA, Ainslie et al., 2023) can be used.”
Competition-level code generation with AlphaCode (2022, methodology): “Therefore, we made two architecture decisions: using (1) an encoder-decoder architecture with asymmetric encoder and decoder structures and (2) the multi-query attention setup from Shazeer (2019) which uses one shared attention head for keys and one for values each block.”
GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding (2020, background): “As explained in Shazeer (2019), memory bandwidth limits are important for incremental decoding with Transformer models.”
GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (2023, background): “Autoregressive decoder inference is a severe bottle-neck for Transformer models due to the memory bandwidth overhead from loading decoder weights and all attention keys and values at every decoding step (Shazeer, 2019; Pope et al., 2022; de Jong et al., 2022).”
Efficiently Scaling Transformer Inference (2022, neutral): “An alternative approach, called multiquery attention (Shazeer, 2019; Chowdhery et al., 2022), still emits nheads for the query tensor, but only a single head for the key and value tensors.”
StarCoder: may the source be with you! (2023, methodology): “Both StarCoder models come with a novel combination of architectural features, such as an 8K token context length (Dao et al., 2022), inﬁlling capabilities through Fill-in-the-Middle (FIM; Bavarian et al., 2022), and fast large-batch inference through Multi-Query-Attention (MQA; Shazeer, 2019).”
Fast Inference from Transformers via Speculative Decoding (2022, background): “Numerous techniques try to make inference more efﬁcient for all tokens, e.g. distillation (Hinton et al., 2015), sparciﬁcation (Jaszczur et al., 2021), quantization (Hubara et al., 2016), and architecture modiﬁcation (So et al., 2021; Shazeer, 2019).”
SantaCoder: don't reach for the stars! (2023, neutral): “You can ﬁnd the full results tables of the text2code experiments are Appendix A. Slight drop in performance for MQA We see a small drop in performance for Multi Query Attention (MQA) compared to Multi Head Attention (MHA).”
Accelerating Large Language Model Decoding with Speculative Sampling (2023, background): “For most applications, auto-regressive sampling (ArS) is highly memory bandwidth bound and thus cannot make eﬀective use of modern accelerator hardware (Shazeer, 2019).”
LambdaNetworks: Modeling Long-Range Interactions Without Attention (2021, methodology): “Rather than attempting to closely approximate an attention kernel as is the case with linear attention, we focus on the efficient design of contextual lambda functions and repurpose a multi-query formulation (Shazeer, 2019) to further reduce computational costs.”
FNet: Mixing Tokens with Fourier Transforms (2021, background): “2 The latter memory footprint results are important as empirical studies have shown that Transformer architectures are often memory-bound (Ivanov et al., 2020; Shazeer, 2019).”
General-purpose, long-context autoregressive modeling with Perceiver AR (2022, methodology): “This strategy is described in detail in Appendix E.3. for previously inferred positions is a common technique for improving generation speed (Vaswani et al.; Shazeer, 2019).”
Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads (2024, methodology): “From a system perspective, LLM inference is predominantly memory-bandwidth-bound (Shazeer, 2019; Kim et al., 2023), with the main latency bottleneck stemming from accelerators’ memory bandwidth rather than arithmetic computations.”
Autoregressive Image Generation without Vector Quantization (2024, methodology): “We use kv-cache [44] for efficient inference.”
The Power of Noise: Redefining Retrieval for RAG Systems (2024, positive): “…highlights both the gold document and the correct answer, “Lando Calrissian”, indicating the expected source and content of the accurate response. will vary according to the single experiment, the instruction will always be placed at the beginning of the prompt, and the query always at the end.”
Yi: Open Foundation Models by 01.AI (2024, methodology): “This approach offers substantial reductions of training and inference costs, compared to the original Multi-Head Attention (MHA) [16, 57, 67].”
A Comprehensive Overview of Large Language Models (2023, positive): “It uses the multiquery attention [121] to reduce memory and cache costs.”