GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

TL;DR

This paper introduces grouped-query attention (GQA), a generalization of multi-query attention, and a method to uptrain existing models with minimal additional compute, achieving near multi-head quality with MQA-like speed.

Contribution

It proposes GQA as a flexible attention mechanism and a recipe for uptraining existing models into GQA with minimal extra compute.

Findings

GQA achieves quality close to multi-head attention.

Uptraining with 5% of original compute is effective.

GQA maintains speed similar to MQA.

Abstract

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.