TransformerモデルにおけるQ, K, Vベクトルの理解

現代の人工知能、特に自然言語処理において、TransformerアーキテクチャはGPT、BERT、その他多くのモデルの基盤となっています。このアーキテクチャの核心には、Query（Q）、Key（K）、Value（V）の3つの重要なベクトルに依存するアテンション機構があります。この記事では、これらのベクトルがトークンからどのように導出されるか、プロンプト処理中にどのように使用されるか、そしてマルチヘッドアテンションとオールリデュース操作がどのように機能して最終出力を生成するかを説明します。

単語からトークンと埋め込みへ

テキスト処理の最初のステップは、各単語（またはサブワード）をトークンに変換することです。モデルの語彙には固定のトークンセットが含まれており、それぞれに一意の整数IDがあります。入力文がトークンに分割された後、各トークンは埋め込みと呼ばれる密なベクトルにマッピングされます。この埋め込みは、トークンの意味を捉える高次元の表現（例：512次元や768次元）です。これらの埋め込みはトレーニング中に学習され、埋め込み行列に格納されます。

Query、Key、Valueベクトルの導出

アテンション機構を可能にするために、各トークンの埋め込みは3つの異なるベクトル（Query、Key、Value）に変換されます。これは、埋め込みベクトルを3つの学習済み重み行列（Wq、Wk、Wv）で乗算することによって行われます。これらの行列はモデルのパラメータであり、トレーニング中に更新されます。結果のベクトルは、元の埋め込み次元と比較してより小さな次元（しばしばd_kと呼ばれる）を持ち、効率的な計算に役立ちます。

Queryベクトル（Q）は、トークンがコンテキスト内で「探している」ものを表します。Keyベクトル（K）は、トークンが「ラベル」として提供するものを表します。Valueベクトル（V）は、集約される実際の情報を含んでいます。アテンション機構は、クエリとキーを比較することによって、各トークンが他のトークンにどの程度注意を払うべきかを決定します。

自己アテンション機構

シーケンス内のすべてのトークンに対してQ、K、Vが得られると、自己アテンション機構はアテンションスコアを計算します。特定のトークンについて、そのクエリはシーケンス内のすべてのトークンのキーと乗算（ドット積）されます。結果のスコアは、各トークンの現在のトークンに対する関連性を示します。これらのスコアは、勾配を安定させるためにキーの次元（d_k）の平方根で割ることによってスケーリングされます。次に、ソフトマックス関数を適用してスコアを合計が1になる確率に変換します。最後に、確率を使用してバリューベクトルの加重和を計算し、トークンの新しいコンテキスト認識ベクトルを生成します。

このプロセスはシーケンス内のすべてのトークンに対して並行して実行されるため、非常に効率的です。新しいベクトルは、フィードフォワードニューラルネットワークと追加の層を通過し、入力の豊かな表現を徐々に構築します。

プロンプト処理と層ごとの計算

プロンプト処理（入力の初期エンコーディング）中に、トークンのシーケンス全体が同時に処理されます。埋め込みはQ、K、Vに変換され、すべてのトークンに対して自己アテンションが計算されます。これはTransformerのすべての層（通常は12、24、またはそれ以上）で繰り返されます。各層は表現を洗練し、モデルが長距離の関係、構文、セマンティクスなどの複雑な依存関係を捉えることを可能にします。最後の層の出力は、各トークンのコンテキスト化された埋め込みを提供します。

推論：新しいトークンの追加

テキスト生成（推論）時、モデルは一度に1つのトークンを生成します。新しいトークンについては、その埋め込みのみが計算されます。次に、そのQueryが導出されます。アテンションを計算するために、モデルは以前のすべてのトークンのKeyとValueを必要とします。これらはKVキャッシュと呼ばれる構造に格納され、再計算を回避します。新しいトークンのQueryはすべてのキャッシュされたKeyと乗算され、ソフトマックスが適用され、キャッシュされたValueの加重和が計算されます。これにより、新しいトークンの新しいコンテキストベクトルが生成され、残りの層と最後に分類層（語彙に対するソフトマックス）を通過して次の単語が選択されます。

この増分アプローチは、以前のトークンに対する重い計算が再利用されるため効率的です。KVキャッシュは、各生成ステップの後に新しいトークンのKeyとValueで更新されます。

マルチヘッドアテンションとオールリデュース

異なるタイプの関係（例：構文、意味、参照）を同時に捉えるために、Transformerはマルチヘッドアテンションを使用します。単一のアテンション操作を実行する代わりに、モデルはQ、K、Vベクトルを複数の小さなヘッド（通常は8、12、または16）に分割します。各ヘッドは入力の異なる側面に焦点を当てることを学習します。ヘッドは並行して動作し、それぞれがベクトルの部分空間上で独自のアテンション機構を計算します。

各ヘッドが出力を生成した後、出力は元の次元の単一のベクトルに連結されます。この連結されたベクトルは、出力投影行列（しばしばWoと呼ばれる）で乗算され、すべてのヘッドからの情報を結合します。この最終ステップはオールリデュース操作に相当します。並列計算を統合された表現にマージします。オールリデュースにより、モデルはヘッド全体での分散学習の恩恵を受けながら、次の層に対して単一の一貫した出力を維持できます。

結論

Query、Key、ValueベクトルはTransformerのアテンション機構の基本です。これらにより、モデルはコンテキスト内の各トークンの重要性を動的に重み付けでき、強力な言語理解と生成を可能にします。これらのベクトルの導出、アテンションの計算、およびオールリデュースによるマルチヘッド並列化の使用は、最新のAIモデルを非常に効果的にしているものです。これらの概念を理解することは、大規模言語モデルを扱ったり研究したりするすべての人にとって不可欠です。