GQA : Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (2023)

Abstract: Multi-query attention (MQA), which only uses a single key-value head, drastically speeds up decoder inference. However, MQA can lead to quality degradation, and moreover it may not be desirable to train a separate model just for faster inference. We (1) propose a recipe for uptraining existing multi-head language model checkpoints into models with MQA using 5% of original pre-training compute, and (2) introduce grouped-query attention (GQA), a generalization of multi-query attention which uses an intermediate (more than one, less than number of query heads) number of key-value heads. We show that uptrained GQA achieves quality close to multi-head attention with comparable speed to MQA.

## introduction autoregressive decoder inference is a severe bottleneck for transformer models due to the memory bandwidth overhead from loading decoder weights and all attention keys and values at every decoding step @xcite @xcite . the memory bandwidth from loading keys and values can be sharply reduced through multi-query attention @xcite , which uses multiple query heads but single key and value heads. however, multi-query attention (mqa) can lead to quality degradation and training instability, and it may not be feasible to train separate models optimized for quality and inference. moreover, while some language models already use multiquery attention, such as palm @xcite , many do not, including publicly available language models such as t5 @xcite and llama @xcite . this work contains two contributions for faster inference with large language models. first, we show that language model checkpoints with multihead attention (mha) can be uptrained @xcite to use mqa with a small fraction of original training compute. this presents a cost-effective method to obtain fast multi-query as well as high-quality mha checkpoints. second, we propose grouped-query attention (gqa), an interpolation between multi-head and multi-query attention with single key and value heads per subgroup of query heads. we show that uptrained gqa achieves quality close to multihead attention while being almost as fast as multiquery attention. ## related work this work is focused on achieving a better tradeoff between decoder quality and inference time through reducing the memory bandwidth overhead @xcite cially helpful for long inputs @xcite . @xcite independently developed gqa with public implementation. a number of other methods have been proposed to reduce memory bandwidth overhead from keys and values, as well as parameters. flash attention @xcite structures the attention computation to avoid materializing the quadratic attention scores, reducing memory and speeding up training. quantization @xcite reduces the size of weights and activations, including keys and values, by lowering precision. model distillation @xcite instead reduces model size at a given precision, using data generated from the larger model to finetune the smaller model. layersparse cross-attention (de jong et al., 2022) eliminates most cross-attention layers which make up the primary expense for longer inputs. speculative sampling @xcite ameliorates the memory bandwidth bottleneck by proposing multiple tokens with a smaller model which are then scored in parallel by a larger model. finally, the uptraining procedure we propose is inspired by @xcite , which uptrains standard t5 checkpoints into sparsely activated mixture-of-experts models. ## conclusion language models are expensive for inference primarily due to the memory bandwidth overhead from loading keys and values. multi-query attention reduces this overhead at the cost of decreased model capacity and quality. we propose to convert multi-head attention models to multi-query models with a small fraction of original pre-training compute. moreover, we introduce grouped-query attention, an interpolation of multi-query and multi-head attention that achieves quality close to multi-head at comparable speed to multi-query attention.