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

要約

単一のキーと値のヘッドのみを使用するマルチクエリ アテンション (MQA) により、デコーダーの推論が大幅に高速化されます。
ただし、MQA は品質の低下につながる可能性があり、さらに、推論を高速化するためだけに別のモデルをトレーニングすることは望ましくない場合があります。
我々は、(1) 元の事前トレーニング コンピューティングの 5% を使用して、既存のマルチヘッド言語モデルのチェックポイントを MQA を備えたモデルにアップトレーニングするレシピを提案し、(2) マルチクエリ アテンションの一般化であるグループ化クエリ アテンション (GQA) を導入します。
これは、中間の (1 つ以上、クエリ ヘッドの数より少ない) 数のキー/値ヘッドを使用します。
アップトレーニングされた GQA は、MQA と同等の速度で、マルチヘッド アテンションに近い品質を達成することを示します。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google