construct each group key and value head by meanpooling all the original heads within that group.
An intermediate number of groups leads to an
interpolated model that is higher quality than MQA
but faster than MHA , and, as we will show, represents a favorable trade-off. Going from MHA
toMQA reducesHkey and value heads to a single key and value head, reducing the size of the
key-value cache and therefore amount of data that
needs to be loaded by a factor of H. However,
larger models generally scale the number of heads,
such that multi-query attention represents a more
aggressive cut in both memory bandwidth and capacity. GQA lets us keep the same proportional
decrease in bandwidth and capacity as model size
increases.
Moreover, larger models suffer relatively less
from memory bandwidth overhead from attention,
as the KV-cache scales with model dimension
while model FLOPs and parameters scale with the
square of model dimension. Finally, standard sharding for large models replicates the single key and
value head by the number of model partitions (Popeet al., 2022); GQA removes the waste from such
partitioning. Therefore, we expect GQA to present
a particularly good trade-off for larger models.
3 Experiments
3.1 Experimental setup
Conﬁgurations All models are based on the
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