a particularly good trade-off for larger models.
3 Experiments
3.1 Experimental setup
Conﬁgurations All models are based on the
T5.1.1 architecture (Raffel et al., 2020), implemented with JAX (Bradbury et al., 2018),
Flax (Heek et al., 2020), and Flaxformer1. For
our main experiments we consider T5 Large and
XXL with multi-head attention, as well as uptrained
versions of T5 XXL with multi-query and groupedquery attention. We apply MQA andGQA to
decoder self-attention and cross-attention, but not
encoder self-attention.
Uptraining Uptrained models are initialized
from public T5.1.1 checkpoints. The key and value
heads are mean-pooled to the appropriate MQA or
GQA structure, and then pre-trained for a further
proportion of original pre-training steps with the
original pre-training setup (Raffel et al., 2020).
Data We evaluate on summarization datasets
CNN/Daily Mail (Nallapati et al., 2016), arXiv
and PubMed (Cohan et al., 2018), MediaSum (Zhu
et al., 2021), and Multi-News (Fabbri et al., 2019);
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