マルチクエリ・アテンションは、大規模言語モデルの推論高速化とメモリ削減を実現する注意機構の改良手法です。複数のクエリヘッド間でキーとバリューを共有することで、メモリ帯域幅のボトルネックを緩和します。
マルチクエリ・アテンションとは
一言でいうと、マルチクエリ・アテンションとは、Transformerモデルのデコーダにおけるキーとバリューのヘッド数を削減し、メモリ効率を劇的に高めるアテンション機構のことです。
詳しく解説
通常のマルチヘッド・アテンションでは、クエリ、キー、バリューのすべてが複数のヘッドに分割され、それぞれ独立して計算されます。しかし、長文の生成処理では、このキーとバリューをメモリにキャッシュし続けることが大きなボトルネックとなります。マルチクエリ・アテンションでは、キーとバリューのヘッドを全クエリヘッドで共通の1つ(または少数)に統合します。これにより、KVキャッシュのメモリ使用量を大幅に削減し、特に推論時のスループットを向上させることが可能です。
具体例・使われ方
超長文のコンテキストを処理する大規模言語モデルや、リアルタイムでの応答が求められる対話型AIシステムの高速化において利用されています。メモリ使用量が削減されるため、限られたGPUメモリ上でもより大きなバッチサイズでの処理が可能になります。
似た用語との違い
従来のマルチヘッド・アテンションと比較して、モデルのパラメータ数や計算精度をほとんど落とさずに、KVキャッシュのサイズを劇的に小さくできる点が異なります。また、一部のヘッドのみを共有するグループトランスファー・アテンションとも混同されやすいですが、マルチクエリ・アテンションは共有を極限まで進めた特例ケースと言えます。
注意点
キーとバリューを共有するため、モデルの表現力がわずかに低下するリスクがあります。そのため、学習初期からの設計やタスクによっては、標準的なマルチヘッド・アテンションと比較して精度に微小な差が生じる場合があるため注意が必要です。