forMHA T5-Large and T5-XXL, and uptrained
MQA andGQA -8XXL models with uptraining
proportion= 0:05. We see that a larger uptrained MQA model provides a favorable tradeoff relative to MHA models, with higher quality
and faster inference than MHA -Large. Moreover,
GQA achieves signiﬁcant additional quality gains,
achieving performance close to MHA -XXL with
speed close to MQA . Table 1 contains full results
for all datasets.
3.3 Ablations
This section presents experiments to investigate
the effect of different modeling choices. We evaluate performance on a representive subsample of
tasks: CNN/Daily Mail, (short-form summarization), MultiNews (long-form summarization), and
TriviaQA (question-answering).
54:454:654:855 55:255:455:6Mean
First
Random
Figure 4: Performance comparison of different checkpoint conversion methods for T5-Large uptrained to
MQA with proportion = 0:05. ‘Mean’ mean-pools
key and value heads, ‘First’ selects the ﬁrst head and
‘Random’ initializes heads from scratch.
Checkpoint conversion Figure 4 compares 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