2048 and output length 512. Going from 1 (MQA) to
8 groups adds modest inference overhead, with increasing cost to adding more groups.
tion (Dao et al., 2022) structures the attention computation to avoid materializing the quadratic attention scores, reducing memory and speeding up
training. Quantization (Dettmers et al., 2022; Zeng
et al., 2022; Frantar et al., 2022) reduces the size of
weights and activations, including keys and values,
by lowering precision. Model distillation (Hinton et al., 2015; Gou et al., 2021) instead reduces
model size at a given precision, using data generated from the larger model to ﬁnetune the smaller
model. Layer-sparse cross-attention (de Jong et al.,
2022) eliminates most cross-attention layers which
make up the primary expense for longer inputs.
Speculative sampling (Chen et al., 2023; Leviathan
et al., 2022; Kim et al., 2023) 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 Komatsuzaki et al. (2022), which
uptrains standard T5 checkpoints into sparsely activated Mixture-of-Experts models.
5 Conclusion
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.
Joshua Ainslie James Lee-Thorp Michiel de Jong Yury Zemlyanskiy Federico Lebrón Sumit Sanghai