Grouped Query Attention (GQA)

Grouped Query Attention (GQA) は、複数のクエリヘッドが、より少ないキーヘッドおよびバリューヘッドのセットを共有するアテンション設計です。マルチヘッドクエリの多様な視点を維持しながら、キーおよびバリューに必要なメモリとデータ移動を削減します。GQAは、モデルが1回に1トークンを生成し、保存されたアテンション状態を繰り返し読み取る自己回帰推論において特に有用です。

GQAの仕組み#

アテンション機構内では、各トークンが3つの表現に射影されます。

• クエリ (Query): 現在のトークンが求めている情報を記述します。
• キー (Key): 利用可能な各トークンが何を表すかを記述します。
• バリュー (Value): クエリがキーに一致したときに取得される情報を含みます。

従来のセルフアテンションでは、モデルが異なる関係性を並行して学習できるように、これらの射影がヘッドに分割されます。GQAでは多数のクエリヘッドを維持しつつ、それらのグループを共有のキーおよびバリューヘッドに割り当てます。例えば、アテンションレイヤーで8個のクエリヘッドと2個のキー・バリューヘッドを使用する場合があります。この場合、各キー・バリューヘッドが4個のクエリヘッドにサービスを提供します。

アテンションの計算自体はスケルドット積アテンション（scaled dot-product attention）のままです。変化するのは、ヘッドの配置と、生成、保存、読み取りが行われるキー・バリューデータの量です。したがって、PyTorch scaled dot-product attentionなどのフレームワークでは、GQAが有効な場合、クエリヘッドの数がキー・バリューヘッドの数で割り切れる必要があります。(docs.pytorch.org)

GQAとマルチヘッドおよびマルチクエリ・アテンションの比較#

GQAは、マルチヘッドアテンションとマルチクエリ・アテンションの中間に位置します。

• マルチヘッドアテンション: すべてのクエリヘッドが独自のキーヘッドとバリューヘッドを持ちます。これにより最大のヘッド独立性が得られますが、最大のキー・バリューメモリ要件が発生します。
• グループド・クエリ・アテンション: 複数のクエリヘッドが各キー・バリューヘッドを共有します。表現力とメモリ効率のバランスを取ります。
• マルチクエリ・アテンション: すべてのクエリヘッドが1つのキーヘッドと1つのバリューヘッドを共有します。これによりキー・バリューのストレージが最小化されますが、キー・バリューの多様性が低下します。

レイヤーに32個のクエリヘッドがある場合、マルチヘッドアテンションでは32個のキー・バリューヘッドを使用する可能性があり、GQAでは8個、マルチクエリ・アテンションでは1個を使用します。GQAは、より広範なTransformerアーキテクチャの代替ではなく、Transformerアテンションレイヤー内の1つの可能な設定です。NVIDIAのmulti-head, multi-query, and grouped-query attention documentationでは、これらのバリエーションの主な違いは、いくつキー・バリューヘッドがクエリヘッドにサービスを提供するかにあると説明されています。(nvidia.github.io)

GQAが重要な理由#

自己回帰生成の間、モデルは以前のキーとバリューをKVキャッシュに保存します。標準的なマルチヘッドアテンションでは、すべてのアテンションヘッドが個別のキャッシュされたキーとバリューを提供します。したがって、長いシーケンス、大きなバッチ、および多くのモデルレイヤーによって、大量のGPUメモリが消費される可能性があります。

GQAはより少ないキー・バリューヘッドを使用するため、そのキャッシュはそれらのヘッドの削減に比例して小さくなります。これにより以下が改善されます。

• メモリ容量: より長いプロンプトや、より多くの同時リクエストを利用可能なメモリに収めることができます。
• 生成スループット: 生成されるすべてのトークンに対して読み取る必要のあるキー・バリューデータが少なくなります。
• 推論レイテンシ: メモリトラフィックの削減により、トークン間の処理時間を短縮できます。
• コンテキストウィンドウのスケーラビリティ: 長いシーケンスの処理がより実用的になります。

これらの向上は、ハードウェア、シーケンス長、バッチサイズ、およびカーネルのサポートに依存します。本番環境のランタイムでは依然として効率的なキャッシュ割り当てが必要であり、例えばNVIDIAのTensorRT-LLM KV cache systemは、GQAのサポートと、ブロックベースのキャッシュ、再利用、オフロードを組み合わせています。(nvidia.github.io)

実社会での応用#

長期コンテキストアシスタント: コーディングアシスタントは、回答を生成する前に数千のソースコードトークンを処理する場合があります。GQAはその履歴に関連付けられたキャッシュ済みキー・バリュー状態を削減し、GPUメモリを比例して増加させることなく、サーバーがより長いファイルやより多くの同時ユーザーを処理できるようにします。

マルチモーダル画像およびビデオアシスタント: ビジョン言語モデルは、画像パッチやビデオフレームを長いトークンシーケンスとして表現できます。GQAは、これらのビジュアルトークンと言語デコーダーに入った後のアテンションキャッシュの圧迫を軽減し、追加の画像、フレーム、またはリクエストにより多くのメモリを残す可能性があります。これは、自己回帰的なKVキャッシュなしでアテンションがすべての画像パッチを並行して処理する標準的なVision Transformerとは異なります。

実践的な実装とトレードオフ#

PyTorchは enable_gqa=True を通じてGQAを公開しています。この文書化されたCUDAの例では、8個のクエリヘッドと2個の共有キー・バリューヘッドを使用しています。

import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

assert torch.cuda.is_available(), "A CUDA device is required."

query = torch.randn(1, 8, 16, 32, device="cuda")
key = torch.randn(1, 2, 16, 32, device="cuda")
value = torch.randn(1, 2, 16, 32, device="cuda")

# Four query heads share each key-value head.
with sdpa_kernel(SDPBackend.MATH):
    output = F.scaled_dot_product_attention(
        query,
        key,
        value,
        enable_gqa=True,
    )

print(output.shape)

出力は8個のクエリヘッドを維持し、キーとバリューは2個のヘッドのみを使用します。PyTorch attention backend selectorはどのサポートされているカーネルが演算を実行するかを制御でき、一方PyTorch Transformer building-block guideはより広範な実装のコンテキストを提供します。(docs.pytorch.org)

GQAはモデルの設計時または適応時に選択する必要があり、既存のチェックポイントに対する推論時の切り替えスイッチではありません。開発者は、フレームワークのサポート、ヘッドの割り切れやすさ、数値的挙動、メモリ使用量、およびタスクの品質を確認する必要があります。コンピュータビジョンの導入においては、Ultralytics TensorRT integrationやUltralytics benchmark modeなどの補完的なワークフローが、アーキテクチャおよびランタイムの最適化がターゲットハードウェア上で意味のある改善をもたらすかどうかを測定するのに役立ちます。