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.
核心贡献 · Key contributions
提出一种方法，仅用原始预训练算力的 5%将多头语言模型检查点升级为多查询注意力模型。 Proposes a recipe to uptrain multi-head language model checkpoints into multi-query attention models using only 5% of original pre-training compute.
引入分组查询注意力（GQA），这是多查询注意力的一种泛化，使用中间数量的键值头。 Introduces grouped-query attention (GQA), a generalization of multi-query attention with an intermediate number of key-value heads.
表明升级后的 GQA 在质量上接近多头注意力，同时速度几乎与多查询注意力相当。 Shows that uptrained GQA achieves quality close to multi-head attention while being almost as fast as multi-query attention.
证明在升级训练中，对键头和值头进行均值池化比其他转换方法效果更好。 Demonstrates that mean-pooling key and value heads works better than other conversion methods for uptraining.
在摘要、翻译和问答任务上提供实证评估。 Provides empirical evaluation on summarization, translation, and question answering tasks.
讨论大型语言模型在质量和推理速度之间的权衡。 Discusses the trade-off between quality and inference speed for large language models.
局限 · Limitations
摘要评估依赖 Rouge 分数，该指标存在缺陷。 Evaluation relies on Rouge score for summarization, which is a flawed metric.
由于算力限制，未将 XXL GQA 模型与从头训练的模型进行比较。 Does not compare XXL GQA model to a model trained from scratch due to limited computation.
仅在编码器-解码器模型上评估，未涉及仅解码器模型。 Only evaluates on encoder-decoder models, not decoder-only models.
质量权衡可能无法推广到更长序列或其他任务。 The quality trade-off may not generalize to longer sequences or other tasks.
升级训练方法可能并非对所有模型架构或规模最优。 Uptraining recipe may not be optimal for all model architectures or sizes.