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.
2 Method
2.1 Uptraining
Generating a multi-query model from a multi-head
model takes place in two steps: ﬁrst, converting the
checkpoint, and second, additional pre-training to
allow the model to adapt to its new structure. Figure 1 shows the process for converting a multi-head
checkpoint into a multi-query checkpoint. The projection matrices for key and value heads are mean
pooled into single projection matrices, which we
ﬁnd works better than selecting a single key and
value head or randomly initializing new key and
value heads from scratch.
Figure 1: Overview of conversion from multi-head to
multi-query attention. Key and value projection matrices from all heads are mean pooled into a single head.
The converted checkpoint is then pre-trained forarXiv:2305.13245v1  [cs.CL]  22 May 2023
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