GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, Joshua Ainslie+, N/A, arXiv'23

[AkihikoWatanabe]

URL

• https://arxiv.org/abs/2305.13245

Affiliations

• Joshua Ainslie, N/A
• James Lee-Thorp, N/A
• Michiel de Jong, N/A
• Yury Zemlyanskiy, N/A
• Federico Lebrón, N/A
• Sumit Sanghai, N/A

Abstract

• Multi-query attention (MQA), which only uses a single key-value head,drastically speeds up decoder inference. However, MQA can lead to qualitydegradation, and moreover it may not be desirable to train a separate modeljust for faster inference. We (1) propose a recipe for uptraining existingmulti-head language model checkpoints into models with MQA using 5% of originalpre-training compute, and (2) introduce grouped-query attention (GQA), ageneralization of multi-query attention which uses an intermediate (more thanone, less than number of query heads) number of key-value heads. We show thatuptrained GQA achieves quality close to multi-head attention with comparablespeed to MQA.

Translation (by gpt-3.5-turbo)

• Multi-query attention（MQA）は、単一のkey-value headのみを使用しており、デコーダーの推論を劇的に高速化しています。
  ただし、MQAは品質の低下を引き起こす可能性があり、さらには、より速い推論のためだけに別個のモデルをトレーニングすることが望ましくない場合もあります。
  私たちは、既存のマルチヘッド言語モデルのチェックポイントを、オリジナルの事前トレーニング計算量の5%を使用してMQAを持つモデルにアップトレーニングするためのレシピを提案し、
  さらに、複数のkey-value headを使用するマルチクエリアテンションの一般化であるグループ化クエリアテンション（GQA）を紹介します。
  私たちは、アップトレーニングされたGQAが、MQAと同等の速度でマルチヘッドアテンションに匹敵する品質を達成することを示しています。

Summary (by gpt-3.5-turbo)

• Multi-query attention（MQA）は、単一のkey-value headのみを使用しており、デコーダーの推論を劇的に高速化しています。
  ただし、MQAは品質の低下を引き起こす可能性があり、さらには、より速い推論のためだけに別個のモデルをトレーニングすることが望ましくない場合もあります。
  既存のマルチヘッド言語モデルのチェックポイントを、オリジナルの事前トレーニング計量の5%を使用してMQAを持つモデルにアップトレーニングするためのレシピを提案し、
  さらに、複数のkey-value headを使用するマルチクエリアテンションの一般化であるグループ化クエリアテンション（GQA）を紹介します。
  アップトレーニングされたGQAが、MQAと同等の速度でマルチヘッドアテンションに匹敵する品質を達成することを示しています。

[AkihikoWatanabe]

通常のMulti-Head AttentionがQKVが1対1対応なのに対し、Multi Query Attention (MQA) #1272 は全てのQに対してKVを共有する。一方、GQAはグループごとにKVを共有する点で異なる。MQAは大幅にInfeerence` speedが改善するが、精度が劣化する問題があった。この研究では通常のMulti-Head Attentionに対して、オリジナルの事前学習に対して追加の5%の計算量でGQAモデルを学習する手法を提案している。

[AkihikoWatanabe]

Main Result. Multi-Head Attentionに対して、inference timeが大幅に改善しているが、Multi-Query Attentionよりも高い性能を維持している。