[論文レビュー] GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

Joshua Ainslie, James Lee-Thorp+4 more|arXiv (Cornell University)|May 22, 2023

Topic Modeling被引用数 27

ひとこと要約

この論文は uptraining を導入し、マルチヘッド・アテンションをマルチクエリと一般化された grouped-query attention（GQA）に変換することで、元の事前学習計算のごく一部を使用しつつ、競争力の質を保ちながら高速推論を実現します。

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.

研究の動機と目的

• autoregressive transformers のメモリ帯域幅に起因するデコーダのボトルネックを動機づけ、品質を損なうことなく推論を高速化する。
• 既存のマルチヘッドチェックポイントをマルチクエリ構成へ変換するコスト効果の高い uptraining レシピを提案する。
• grouped-query attention（GQA）を、マルチヘッドとマルチクエリの間の補間として導入する。
• uptrained GQA が、タスクを跨いでマルチヘッド・アテンションに近い品質を、マルチクエリ・アテンションに近い速度で達成することを示す。

提案手法

• マルチヘッド・アテンションのチェックポイントを、キー・値射影行列を平均プーリングして単一ヘッドを形成することで MQA 用、または各 GQA グループ用に変換する。
• 同じデータとレシピを用い、元の事前学習ステップのごく一部（α）だけで変換後モデルを uptrain する。
• GQA を、G 個のクエリ・ヘッドのグループで定義し、各グループにつき1つの KV ヘッドを共有し、MQA（G=1）と MHA（G=H）の間を補間する。
• デコーダ自己注意とクロス注意に対して、T5-Large および T5-XXL を用い、エンコーダ自己注意には適用しない形で MQA および GQA を実験する。
• 要約、翻訳、QA ベンチマークで推論時間と dev 集の性能を報告して評価する。

実験結果

リサーチクエスチョン

• RQ1 マルチヘッド・チェックポイントは、追加計算を限定した上で効果的に高速なマルチクエリ形へ uptraining できるか。
• RQ2 grouped-query attention は速度と品質の間の調整可能なトレードオフを提供し、MQA を超えつつ MHA に近づくか。
• RQ3 uptraining は異なるタスクとモデル規模での安定性と性能にどのように影響するか。

主な発見

• uptrained MQA は MHA-Large より推論が速く、品質も高いが、場合によっては XXL-MHA に劣ることもある。
• uptrained GQA は MHA-XXL に近い品質を、MQA に近い速度で実現し、好ましいトレードオフを提供する。
• チェックポイント変換の平均プーリングは、ヘッドを選択するまたはランダム初期化を行うより情報を良く保持する。
• 5% の uptraining で性能が向上し、MQA と GQA の両方で 5–10% 超えると収益が減少する。
• GQA のグループ数が増えるほど帯域幅ダイナミクスとキャッシュを考慮したメリットが大きくなる一方、コストは増大する。
• グループ数を 1（MQA）から増やすと、速度向上は modest、コスト増が顕著になる。8 グループが有利な中庸と特定された。