GQA-8-XXL 0.28 47.1 43.5 45.4 47.7 36.3 47.2 28.4 81.6
Table 1: Inference time and average dev set performance comparison of T5 Large and XXL models with multihead attention, and 5% uptrained T5-XXL models with multi-query and grouped-query attention on summarization
datasets CNN/Daily Mail, arXiv, PubMed, MediaSum, and MultiNews, translation dataset WMT, and questionanswering dataset TriviaQA.
Fine-tuning For ﬁne-tuning, we use a constant
learning rate of 0.001, batch size 128, and dropout
rate 0.1 for all tasks. CNN/Daily Mail and WMT
use input length of 512 and output length 256.
Other summarization datasets use input length
2048 and output length 512. Finally, TriviaQA
uses input length 2048 and output length 32. We
train until convergence and select the checkpoint
with the highest dev performance. We use greedy
decoding for inference.
Timing We report time per sample per TPUv4
chip, as measured by xprof (Google, 2020). For
timing experiments we use 8 TPUs with the largest
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