グループ化クエリー注意法（GQA）とは

Ainslieらは、2017年の画期的な論文「Attention is All You Need」で導入された革新的な自己注意アルゴリズムであるマルチヘッド注意法（MHA）の最適化として、グループ化クエリー注意法を考案しました。この論文は、Transformerニューラル・ネットワークを確立しました。より具体的には、MHAの初期の最適化であるマルチクエリー注意法（MQA）の一般化とより制限されたアプリケーションとして提案されました。

標準的なマルチヘッド注意法は、機械学習、自然言語処理（NLP）、および生成AIを比較的に進化させましたが、計算リソースとメモリー帯域幅に対する要求が非常に高くなります。LLMがより大規模かつ高度になるにつれ、これらのメモリー使用量要件が進歩のボトルネックになりました。特に、テキスト生成、要約、およびその他の生成AIタスクに使用される自動回帰デコーダーのみのLLMではボトルネックになりました。

その後の研究は、マルチヘッド注意法を強化または合理化する手法に重点が置かれました。フラッシュ注意法やリング注意法などの一部の研究では、モデルのトレーニングと実行に使用されるGPUが計算とメモリー・ストレージを処理する方法が改善されました。GQAやMQAなどの他の研究では、Transformerアーキテクチャーがトークンを処理する方法の変更が調査されました。

グループ化クエリー注意法は、標準のマルチヘッド注意法とマルチクエリー注意法のトレードオフのバランスを取ることを目的としています。前者は、メモリー帯域幅のオーバーヘッドの増加と速度の低下を犠牲にして、精度を最大化します。後者は、精度を犠牲にして、速度と効率を最大化します。

グループ化クエリー注意法がTransformerモデルを最適化する仕組みを理解するには、まずマルチヘッド注意法が一般的にどのように機能するかを理解することが重要です。GQAとMQAはどちらも、MHAのコア・メソッドを置き換えるのではなく、改良したものに過ぎません。

LLMやTransformerアーキテクチャーを使用するその他のモデルの原動力は、シーケンス内の異なるトークン間の関係を理解するための数学的フレームワークである自己注意法です。自己注意法により、LLMは静的なベースライン定義だけでなく、他の単語やフレーズによって提供されるコンテキストを通じてテキスト・データを解釈できます。

テキスト生成に使用される自己回帰LLMでは、注目メカニズムによって、その時点でどの前のトークンに最も「注目する」価値があるかを判断することで、モデルがシーケンス内の次のトークンを予測するのに役立ちます。最も関連性が高いと判断されたトークンの情報には、より大きな注目重みが与えられ、関連性がないと判断されたトークンの情報には、0に近い注意重みが与えられます。

Transformerモデルをアニメーション化するマルチヘッド注意メカニズムは、注意層を複数の注意ヘッドに分割して、自己注意を何度も並列に計算することにより、豊富なコンテキスト情報を生成します。

影響力のある論文「Attention is All You Need」では、リレーショナル・データベースの用語であるクエリー、キー、値を使用して、その注意メカニズムを明確に説明しました。リレーショナル・データベースは、関連データの保存と取得を簡素化するように設計されています。各データに一意の識別子（「キー」）を割り当て、各キーは対応する値に関連付けられます。リレーショナル・データベースの目的は、各クエリーを適切なキーと一致させることです。

シーケンス内の各トークンに対して、マルチヘッド注意法では3つのベクトルを作成する必要があります。

• トークンが「求めている」情報を表すクエリー・ベクトル、Q。例えば、名詞のクエリーベクトルは、それを記述する形容詞の検索を表す場合があります。
• トークンに含まれる情報を表すキー・ベクトル、K。各トークンのキー・ベクトルと他の各トークンのクエリー・ベクトルとの関連性を表すアライメント・スコアは、注目度の重みを計算するために使用されます。
• 他のトークンのキーベクトルからの注目度重み付けされた寄与によって更新されるコンテキスト情報を表す値ベクトル、V。

注意メカニズムによって媒介されるこれら3つのベクトル間の数学的相互作用は、モデルが各トークンのコンテキスト固有の理解を調整する方法です。
 

特定のトークンに対してこれら3つのベクトルを生成するために、モデルは、そのトークンの元のベクトル埋め込みから開始します。ベクトルの各次元がトークンの意味の抽象要素に対応する数値エンコードです。これらのベクトルの次元数は、事前に決定されたハイパーパラメーターです。

各トークンのQ、K、Vベクトルは、最初の注意層の前にある線形層に元のトークン埋め込みを渡すことによって生成されます。この線形層は、モデル重みの3つの一意のマトリックス（W_Q、W_K、W_V）に分割されます。その中の特定の重み値は、テキスト例の膨大なデータセットに対する自己教師あり事前トレーニングを通じて学習されます。

トークンの元のベクトル埋め込みにW_Q、W_K、W_Vを掛けると、それぞれ対応するクエリー・ベクトル、キー・ベクトル、値ベクトルが生成されます。各ベクトルに含まれる次元数dは、各重み行列のサイズによって決まります。QとKの次元数は同じで、d_kになります。

これらの3つのベクトルは、注意層に渡されます。

注意層では、Q、K、およびVベクトルを使用して、シーケンス内の各位置にある各トークン間のアライメント・スコアを計算します。次に、これらのアライメント・スコアは、ソフトマックス関数を使用して注意重みに正規化されます。

シーケンス内の各トークンxについて、そのトークンのクエリー・ベクトルQ_xと他の各トークンのキー・ベクトルKのドット積を計算することによって、アライメント・スコアが計算されます。言い換えると、これらを掛け合わせることです。2つのトークン間の意味のある関係がそれぞれのベクトル間の類似性に反映されている場合、これらを掛け合わせると大きな値が得られます。2つのベクトルがアライメントされていない場合、これらを掛け合わせると小さな値または負の値が得られます。ほとんどのTransformerモデルでは、拡張ドット積注意法と呼ばれるバリエーションが使用されています。このバリエーションでは、QKがスケール、つまり$\frac{1}{\sqrt{d_K}}$倍されて、トレーニングの安定性が向上します。

これらのクエリー・キーのアラインメント・スコアは、ソフトマックス関数に入力されます。ソフトマックスは、すべての入力を0から1の間の値に正規化し、すべての入力を合計すると1になります。ソフトマックス関数の出力は、注意重みです。各注意重みは、トークンxの注意が他のトークンに支払われる割合（1のうち）を表します。トークンの注意重みが0に近い場合、そのトークンは無視されます。注意重みが1の場合、トークンはxの注意をすべて受け取り、他のトークンはすべて無視されます。

最後に、各トークンの値ベクトルにその注目度の重みが掛けられます。以前の各トークンからの注目度の重み付けされた寄与は平均化され、トークンxの元のベクトル埋め込みに追加されます。これにより、トークンxの埋め込みは、シーケンス内の関連する他のトークンによって提供されるコンテキストを反映するように更新されます。

更新されたベクトル埋め込みは、独自の重み行列W _Zを持つ別の線形層に送られ、そこでコンテキスト更新されたベクトルが一貫した次元数に正規化されてから、次の注意層に送られます。各漸進的な注意層は、より大きなコンテキストのニュアンスを捉えます。

各注意重み付けコンテキストを個別に考慮するのではなく、他のトークンからの注目度重み付け貢献の平均を使用すると、数学的には効率的ですが、詳細が失われます。

これを補うために、Transformerネットワークは元の入力トークンの埋め込みをh個の均等なサイズの部分に分割します。同様に、W_Q、W_K、およびW_Vをそれぞれクエリー・ヘッド、キー・ヘッド、および値ヘッドと呼ばれるh個のサブセットに分割します。各クエリー・ヘッド、キー・ヘッド、および値ヘッドは、元のトークン埋め込みの一部を受け取ります。これらのクエリー・ヘッド、キー・ヘッド、および値ヘッドの並列トリプレットのそれぞれによって生成されたベクトルは、対応する注意ヘッドに送られます。最終的に、これらのh個の並列回路の出力が連結され、完全なトークン埋め込みが更新されます。

トレーニングでは、各回路は意味の個別の側面を捉える個別の重みを学習します。これにより、モデルは単語の意味がその周囲の単語のコンテキストによってどのように影響を受けるかをさまざまな方法で処理できるようになります。

標準マルチヘッド注意法の欠点は、重大な欠陥が存在することではなく、最適化されていないことです。MHAはこの種のアルゴリズムとしては初めてのもので、注意計算の一般的なメカニズムを最も複雑に実行しています。

MHAの非効率性の大部分は、計算とモデル・パラメーターの多さに起因します。標準的なMHAでは、各注意ブロック内の各クエリー・ヘッド、キー・ヘッド、および値ヘッドには、それぞれ独自の重み行列があります。そのため、例えば、各注意層に8つの注意ヘッド（ほとんどの最新のLLMよりはるかに少ない）があるモデルでは、層のQ、K、およびVヘッドだけで24個の一意の重み行列が必要になります。これにより、各層で膨大な数の中間計算が必要になります。

この構成の結果、計算コストが高くなります。MHAの計算要件は、シーケンスの長さに対して2乗で増加します。つまり、入力シーケンスのトークン数が2倍になると、複雑さは4倍になります。これにより、コンテキスト・ウィンドウのサイズに実際的な制限が課せられます。

MHAはシステム・メモリーにも大きな負担をかけます。GPUには、後続の各処理ステップで呼び出す必要のある膨大な量の中間計算の出力を保存するためのオンボード・メモリーがあまりありません。これらの中間結果は、代わりにGPUチップ自体には配置されていない高帯域幅メモリー（HBM）に保存されます。これにより、キーと値をメモリーから読み取るたびに、わずかな遅延が発生します。Transformerモデルが数十億のパラメーターに拡張され始めると、トレーニングと推論の実行に必要な時間と計算がモデル・パフォーマンスのボトルネックになりました。

さらなる進歩には、複雑な言語パターンを学習して再現するTransformerの能力を低下させることなく、計算ステップの数を減らす方法が必要でした。このような状況で、MQA、そしてその後のGQAが導入されました。

マルチクエリー注意法（MQA）は、マルチヘッド注意法を簡素化してメモリー使用量と中間計算を削減する、計算効率の高い注意メカニズムです。各注意ヘッドに固有のキー・ヘッドと値ヘッドをトレーニングする代わりに、MQAは各層で単一のキー・ヘッドと単一の値ヘッドを使用します。したがって、キー・ベクトルと値ベクトルは1回だけ計算され、この1セットのキー・ベクトルと値ベクトルは、すべてのh注意ヘッド間で共有されます。

この簡素化により、モデルが計算して高帯域幅メモリーに保存する必要がある線形投影の数が大幅に削減されます。MQAを紹介した2019年の論文によると、MQAを使用すると、キーと値のペアのストレージ（またはKVキャッシュ）が10～100倍小さくなり、デコーダーの推論が12倍高速になります。MQAのメモリー使用量の削減により、バッチ・サイズを大きくできるため、トレーニングも大幅に高速化されます。

MQAにはメリットがある一方で、避けられない欠点もいくつかあります。

• 性能の低下：当然のことながら、トレーニング可能な固有のモデル・パラメーターの数を減らすと、モデルの知識やニュアンスに対する能力が低下します。MQAは、標準のMHAと比較して精度が大幅に低下するため、特定の状況やユースケースには適していません。
  
  
• 最初からトレーニングする必要あり：標準MHAでトレーニングされたモデルは、MQAに単純に適応させることはできません。代わりに、最初からMQAを使用してトレーニングする必要があります。つまり、MQAを使用して既存のモデルを最適化することはできず、新しいモデルでMQAを試す際にはかなりの機会費用がかかります。
  
  
• テンソル並列処理の冗長性：GPUでTransformerモデルをトレーニングする主なメリットは、複数の複雑なテンソル演算を並列に実行できることです。K値とV値は、これらの操作を実行するGPUクラスターの各ノードに存在する必要があります。つまり、実際には各ノードに対して複製される必要があります。これは、標準のMHAよりも効率的ではありますが、コンピューティング・リソースの最適な使用方法ではありません。

グループ化クエリー注意法は、マルチクエリー注意法のより一般的で柔軟な定式化であり、すべてのクエリー・ヘッドで1つのキーと値のセットを共有するのではなく、クエリー・ヘッドを複数のグループに分割し、各グループでキーと値のセットを共有します。

2023年5月に「GQA：マルチヘッド・チェックポイントからの一般化マルチクエリーTransformerモデルのトレーニング」が公開されて以来、多くのLLMがすぐにGQAを採用しました。例えば、Meta社は2023年7月にLlama 2モデルに初めてGQAを採用し、2024年にリリースされたLlama 3 モデルでもGQAを維持していました。Mistral AI社は2023年9月にリリースしたMistral 7BモデルでGQAを使用しました。同様に、IBMのGranite 3.0モデルは、高速推論のためにGQAを採用しています。同様に、IBMのGranite 3.0モデルは高速推論のためにGQAを採用しています。

理論的には、GQAは標準MHAと完全なMQAの間のスペクトルの一般化として考えることができます。注意ヘッドと同じ数のキー値ヘッド・グループを持つGQAは標準MHAと同等であり、1つのヘッド・グループを持つGQAはMQAと同等です。

実際には、GQAはほとんどの場合、グループの数自体が重要なハイパーパラメータとなる中間的なアプローチを意味します。

グループ化クエリー注意法にはいくつかのメリットがあり、主要なLLMで比較的広く採用されています。

• 効率的なGPU使用法：GQAによるキーと値のペアの配布は、テンソル並列処理を利用し、冗長な値を複製することで「無駄」となるコンピューティングの量を削減します。
  
  
• 効果的な妥協点：GQAは、デコーダーの推論速度とパフォーマンスの精度の間で理想的なトレードオフを提供します。つまり、MHAとほぼ同等の精度でありながら、MQAとほぼ同等の速度を実現します。
  
  
• メモリー帯域幅オーバーヘッドの削減：MQAと同様、GQAは推論時に計算、保存、取得が必要な中間計算の数を大幅に削減します。
  
  
• 柔軟なトレーニング：MQAとは異なり、グループ化クエリー注意法では、このアプローチを使用してモデルを最初からトレーニングする必要はありません。標準のMHAを使用して事前トレーニングされたモデルは、「アップトレーニング」と呼ばれるファイン・チューニング・プロセスを通じてGQAを使用するように適応できます。