トランスフォーマーの自己アテンションの理解⑥クエリとキーとバリュー
この記事では、トランスフォーマーのアテンションの仕組みの中心となる概念、クエリ(Query)・キー(Key)・バリュー(Value)を見ていきます。ようやく自己アテンションの解説に手が届きます。
では、さっそく始めましょう。
アテンションの概要
RNNのエンコーダ・デコーダでもアテンションの仕組みがありました。それはエンコーダから出力された隠れ状態(文脈)とデコーダ内の隠れ状態との関係を計算して、エンコーダからどの隠れ状態をより多く取り込むのかを計算する手法でした。また、計算としては、ニューラルネットワーク(関連論文)やベクトルの内積(関連論文)などを使います。
下図では、ベクトルの内積を使っている想定になっています。後に解説しますが、これが内積を使ったアテンションの計算の基本形になります。

トランスフォーマーでも同様な手法を使うのですが、エンコーダとデコーダの間だけでなく、文章内のトークン同士の関係を計算することにも使います。それを自己アテンションと呼びます。
下図は、文章内の各トークンから同じ文章内の全てのトークンに対して関連性の強さの計算をしている様子を表現しています。言葉が入った箱は埋め込みベクトル(位置エンコーディングを含む)になります。

トークン同士の関連の強さを知ると、どのトークンがどのトークンと相対的により強い繋がりがあるのかが分かります。
下図は、実際のデータから生成したものではありませんが、トークン同士の関係の強さを色で表しています。明るい色ほど関係が強くなっています。

対角線上は同じトークンの関係、例えば「犬」から「犬」のトークンへの関係となり、通常は関係が相対的に強く現れます。それ以外でも、「犬」から「歩く」の関係が強く、それを文脈としてより多く取り入れる、つまり「歩く」の埋め込みからの情報を「犬」の埋め込みベクトルへと追加することが望まれます。
このようにアテンションの計算ではトークン同士の関係の強さを知ることが第一歩になります。
ただし、埋め込みベクトルをそのまま使って内積を取るわけではありません。直接の内積ではトークン同士の分散表現がどのくらい似ているのかが分かるのですが、トークンの埋め込みベクトルには様々な意味や文脈が込められているので、すべてを同時に比べると関係性の評価も曖昧になってしまいます。
ここで第三弾で登場した「内積で質問をする」考え方が必要となります。
クエリ
例えば、「犬」と「歩く」に関して、ある特性や性質や特徴などに注目して関係の強さを計算したいとします。つまり、「犬」と「歩く」の埋め込みベクトルそのものではなく、特定の特徴をベクトルとして抽出した上で内積を取ることを行います。
「犬」という言葉の分散表現には様々な特徴が含まれており、全てをそのまま使うと比較が困難です。よって、まずはそこから関係の強さを計算するための特徴量を抽出します。この特徴量ベクトルをクエリ(Query、質問)ベクトルあるいは単にクエリと呼びます。
下図にあるように、行列Qを使って、トークンの埋め込みベクトルから質問に関する特徴のベクトルへと変換します。

具体的には、「犬」の埋め込みベクトルと行列Qの掛け算(積)になります。以下に順を追って説明します。
トランスフォーマーの埋め込みベクトルは512次元でしたが、ここでは次元数を抑えた図で解説します。なお、ベクトルは1次元の横に並んだ数値とします。これを行ベクトルあるいは単にベクトルと呼びます。行列は2次元に配置された値を持っています。

行列の列は縦に数値が並んでおり、ベクトルとして扱えるので列ベクトルと呼びます。よって行列Qでは列ベクトルが5つ並んでいます。「犬」のベクトルと行列Qの列ベクトルとの内積を取ることで、以下に説明する手順でクエリ(質問)のベクトルができます。
まず、「犬」のベクトルと行列Qの1番目の列ベクトルとの内積を計算した値がクエリのベクトルの第一要素になります。

続いて、「犬」のベクトルと行列Qの2番目の列ベクトルとの内積を計算した値がクエリのベクトルの第二要素になります。

この作業を繰り返して「犬」のベクトルと行列Qの全ての列ベクトルとの内積を取ります。実際にはこれらの計算はGPUなどを利用して並列して行われますが、ここではわかりやすさを重視して説明しています。
なお、上図ではトークンが4次元で特徴ベクトルが5次元になっています。そのため、行列Qは4x5の行列になっています。このようにトークンのベクトルの次元数とクエリのベクトルの次元数は必ずしも一致する必要はありません。
このようにして質問(クエリ)のための特徴量をもったベクトルが出来上がります。しかし、この特徴量とは何なのでしょうか。
より具体的なイメージを浮かべられるように、第三弾で扱った映画の特徴に対するユーザーの反応の度合いを計算したことを思い起こしてみてください。例えば、ホラー好きであるとか、ゾンビが好き、コメディは嫌い、SFが好きなど、ユーザーの特徴・特色はある意味「ゾンビのいるSFホラーはありますか」という質問だと捉えられます。それを引き出すのが行列Qの役割です。
つまり、行列Qの各列ベクトルは何らかの特徴・特色を表しています。その特徴・特色の列ベクトルと「犬」の埋め込みベクトルとの内積が大きくなれば、「犬」にはその特徴・特色が強くあることになります。言い方を変えると、行列Qの中にある5つの列ベクトルに代表される特徴・特色に対する「犬」ベクトルの反応がクエリになります。
行列Qの中身は訓練によって調節されるので、上述の映画による例えのように簡単には説明できません。また、次に説明するキーとの関連でトークン間の関係性が決まってくるので単純に特徴・特性とも言い切れません。ニューラルネットワークの最適化を通してさまざまな関係を判断するようになります。例えば、文法的なもの(主語・述語、形容詞、副詞など)も含まれるでしょう。
次に、比べる相手のトークンからキーを抽出する手順を見ていきます。
キー
キー(Key)を抽出する作業は、クエリの場合と同様ですが、異なる行列Kを使います。キーとは何なのかは後にまわして、まずは計算方法を見ていきます。
例えば、「歩く」のベクトルと行列Kの各列ベクトルに対する反応がキーのベクトルになります。

具体的には、「歩く」のベクトルと行列Kとの行列乗算を計算すれば、キーとして使う特徴量を抽出したベクトルが生成されます。

計算の手順はクエリの場合と同じです。
さて、キーとは何なのでしょうか。これはクエリとキーをひと組みのものとして捉えるとわかりやすくなります。次に、この点を解説します。
クエリとキーによる注目度
クエリもキーもトークンのベクトルから抽出された特徴量ベクトルで、これらの内積を計算することでトークン間のある特徴・特性に関した関係の強さを計算します。
クエリとキーの内積を計算するためには、キーのベクトルの次元数はクエリのベクトルの次元数と一致する必要があります。

このように、クエリとキーは別々のものというより、一対で意味を持ちます。
あえて直感的な表現をすると、あるトークンからの質問を各トークンへ投げかけ、どのトークンが強く反応するのかがわかるというわけです。アンケートや問診票の質問と返答のような感じになります。質問がない返答は意味がないので一対で意味を持ちます。
つまり、クエリのための行列Qとキーのための行列Kは独立したものではなく、訓練を通してある特徴・特性に関する関係性の強さを図れるように調節された重みを持った行列になります。
ここでは行列Qと行列Kがどのような質問(クエリ)をしているのか、どのような返答(キー)を取り出しているのかを我々は直接コントロールしていませんが、QとKは同時に重みを調節されるので一対のものとなります。
まとめると、クエリやキーの内容がどのように抽出されるのかは、訓練によって形成されていくのであり、人間が意図して設定したものではありません。訓練を通して翻訳のタスクを何度も繰り返していくうちに、行列QとKの値も最適化されていきます。なので、質問と返答という直感的な表現は例えとして理解の助けにはなりますが、あまり強調して区別する必要もありません。
いずれにせよ、「犬」のクエリと「犬」のキーの内積を計算すれば強い反応が出るでしょう。また「犬」のクエリと「歩く」のキーがどう反応するかは、行列QとKの内容に依存します。
ただし、上記の説明だけだとイメージが湧かないと思うので、仮にこんなことが起きているかもしれないというストーリーとして以下に解説します。
行列Qは「犬」から特徴を取り出します。例えば、「四足歩行」、「走るのが得意」、「散歩が好き」、「空は飛べない」、「海に住んでいない」などの情報を含んだ特徴量のベクトルに変換します。

一方で、行列Kは「歩く」から特徴を取り出します。「地面の上で移動する」、「走るよりスピードが遅い」、「空は飛ばない」、「海で泳がない」、「二足歩行か四足歩行」など。こうして「歩く」のキー(クエリに関する重要な特徴)が抽出されました。

「犬」と「歩く」の文章内における位置関係(前後)も位置エンコーディングが各ベクトルに含まれているで抽出されたベクトルにも含まれるでしょう。もし、「歩く」が「犬」より前だったら「歩く」は無関係かもしれません。あるいは、「歩く」が「犬」を修飾しているのかもしれません。

いずれにせよ、行列QとKで抽出されたクエリとキーのベクトルの内積をとることで、抽出された特徴量における「犬」と「歩く」の関係の強さが測れるわけです。
以上から、行列QとKは1組でその役割を果たしているのが直感的にわかります。
クエリとキーと別々の名前で呼んではいますが、その中身は行列QとKの中の列ベクトルに対する反応を抽出したベクトルです。あまり、質問とかキーとかの名前にこだわる必要はありません。むしろ、ニューラルネットワークがトークン同士の関係性を知る手段として行列QとKが与えられており、この二つの組み合わせによって複雑な関係性を捉えることが可能になっていると捉えるべきでしょう。
そして、クエリとキーの内積の値は、あるトークンが文章内の各トークンにどのくらい注目すべきかを教えてくれます。これによって、関係の強いトークンからの文脈をより多く取り込むことができます。
スケールされた内積とソフトマックス
クエリとキーの内積の値を使って各トークンからの情報をどの程度取り込んで行くのかを決めるのですが、内積で計算した値は大小さまざまです。マイナスの値になることもあります。そのままだと比較が難しいので、標準化をします。
例えば、「犬」から各トークンへ、行列QとKを通しての以下のような数値が計算されたとします。

内積の値はソフトマックス(Softmax)を使って合計が1になる重みとして計算されます。
$$
\text{softmax}\left( [\text{クエリのベクトル}] \cdot [\text{キーのベクトル}] \right)
$$
内積ではマイナスの値が出ることもありますが、ソフトマックスは指数関数を使っているので、マイナスが含まれる場合でも合計1になる重みとして標準化することができます。
$$
\text{softmax}(x_i) = \frac{e^{x_i}}{\sum\limits_{k=1}^n e^{x_k}}
$$
ここで$${x_i}$$はトークン間の関係の強さを表す内積で得た値になります。すべての内積の総和で割っているので割合を計算していることになります。

なお、トランスフォーマーが使う内積の計算は、スケールされた内積(Scaled Dot-Product)であり、内積にある調整を加えます。
$$
\frac{[\text{クエリのベクトル}] \cdot [\text{キーのベクトル}]}{\sqrt{次元数}}
$$
上述の例ではクエリとキーの次元数は5なので$${\sqrt{5}}$$で内積の値を割ることになります。
理由は、次元の数が増えるとベクトルの要素の数が増えるので内積の絶対値(正の値はそのまま、マイナス場合はマイナスをとって正の値にしたもの)が大きくなりがちだからです。次元が多いということは、たくさんの値を足し合わせることになるわけで、絶対値の大きい値が出る可能性が高くなります。
注釈:確率統計における分散の考えを使うと、ランダムな値を持つ複数の変数の分散を足し合わせると分散がより大きくなり平均から離れた値が出やすくなります。よって、分散を小さくするために内積の値を次元数の平方根で割っています。分散で考えると次元数で割っているのに等しくなります。
それに加え、ソフトマックスが指数関数を使うため、大きな値がより大きく誇張されます。大きな値と小さな値との差が極端に開くと、標準化した時に小さな重みがほぼゼロに押しやられてしまいます。

そうなると、同じトークン同士(例えば、「犬」と「犬」)のクエリとキーによる内積の値が標準化された際にほぼ1になってしまい、他のトークンへの注目度がゼロになるのでアテンションの意味がなくなってしまいます。
よって、内積の値を次元数の平方根で割ることで値の格差が開きすぎないように調節しています。以上により、スケールされた内積によるソフトマックスの計算は以下になります。
$$
\text{softmax}\left(\frac{[\text{クエリのベクトル}] \cdot [\text{キーのベクトル}]}{\sqrt{次元数}}\right)
$$
これで、あるトークンから文章内の各トークンへの注目度は、スケールされた内積とソフトマックスによって標準化されました。この注目度に従って、そのトークンの埋め込みベクトルに各トークンの埋め込みベクトルからの文脈を取り込んでいきます。
ここで文脈を取り出すバリューの概念が登場します。
バリュー
バリュー(値、Value)もまたトークンの埋め込みベクトルから抽出された特徴量なのですが、クエリやキーとは異なります。行列Vを使います。

しかし、なぜ行列Vが必要なのでしょうか。
これは、各トークンをマッチさせるための特徴量であるクエリやキーのベクトルはあくまでも関係を測るためのものであり、トークンの文脈を取り出すための行列Vが別に必要だからです。
言い方を変えると、トークン同士の関係の強さを計算するための情報は、トークンの埋め込みベクトルから抽出する文脈の情報と必ずしも同じではないため、別の行列Vを使います。
また、後に解説しますが、バリューのベクトルの次元はトークンの埋め込みベクトルの次元と同じである必要があります。トークンの次元が4ならば、行列Vは4行4列になります。これに対して行列QとKによるクエリとキーのベクトルの次元数には制限がありません。
それでもまだ疑問が残ります。行列Vを通さずに、そのままトークンのベクトルの値を使っても良いのではないでしょうか。
これに関しては、行列Vを挟むことで必要に応じでベクトルを操作する能力を与えておく、と考えると納得できます。行列Vの重みは学習の過程で調節されるので、実際に行列Vがトークンのベクトルをどの程度変換するのかは訓練の結果次第です。何もせずにそのまま値を通すのか、それとも何らかの変換を行うのか、それは学習に任せるということです。結果が良ければ、行列Vが何をしてようが構わないわけです。
さらに、行列QとKとVは同時に最適化されるのでクエリやキーの生成に合わせた情報をVによって抽出することができるわけです。よって、クエリ・キー・バリューはひと組みで考えるべきです。
それに比べ、行列Vを持たない場合は上記のような自由度がなくなります。クエリとキーに関係ない情報が混ざったままになるでしょう。むしろ、様々な関数として機能できるニューラルネットワークによって、行列QとKとVを同時に最適化をさせるのが最善です。
こうして得られたバリューの各ベクトルを異なる色で下図のように表現しておきます。色に関して深い意味はありませんが、元のベクトルから変化している(あるいはしていない)度合いを表現しています。

抽出された文脈であるバリューの値を注目度に従って取り込んでいくことで、自己アテンションの計算が行えます。
自己アテンション
スケールされた内積とソフトマックスによる注目度を、各トークンのバリューへの重みとして使います。
下図では、「犬」のトークンによる各トークンのバリューへの注目度が表現されています。それぞれ、「犬」に95%、「が」に2%、「歩く」に3%になっています。

「犬」が「犬」からの文脈を最も多く取り込むのは妥当でしょう。また、「歩く」の注目度が「が」より少し高いのも自然な感じがします。もちろん、これらの値は行列QとKにも依存するので、何の特徴量を比べているのかがわからないと断言はできないので、イメージとして理解してください。
この注目度に応じてバリューのベクトルからの値を取り込みます。つまり、ここで加重平均をとります。

こうして出来たベクトルが、更新された「犬」のトークンのベクトルとなります。つまり、「犬」の埋め込みベクトルに文章からの文脈を取り込みました。

上記の例では、更新された「犬」の埋め込みベクトルは、「犬」のバリューから多くの値を取り入れています。そして、他のトークンからも重みに応じてバリューを取り入れています。
また、バリューのベクトルの次元数がトークンの埋め込みベクトルの次元数と同じである必要があるのもわかります。
同様な処理を各トークンに施せば、すべてのトークンは関係の強さに応じて文章からの文脈を取り組むことができます。
以上が、トランスフォーマーにおける自己アテンション(Self-Attention)の基本的な仕組みです。クエリとキーによる各トークン同士の関係の強さを重みとしてバリューの加重平均を計算したものが、更新された埋め込みベクトルとなります。このようにして再帰の構造なしに文脈を取り入れているわけです。
ここまで理解できると、その仕組みが割と単純なものであると感じるでしょう。第五弾でも解説しましたが、トランスフォーマーの論文の冒頭においてGoogle Brainの研究者であるAshish Vaswaniらは「新しいシンプルなネットワークアーキテクチャであるトランスフォーマーを提案します」と宣言しました。
まとめ
トランスフォーマーの自己アテンションの仕組みをクエリ、キー、バリューの概念を通して行列QとKとVを使って説明してきました。
トランスフォーマーでは、自己アテンションの他にソース・ターゲットのアテンションがあります。詳しい説明は次回にしますが、自己アテンションの仕組みの応用として考えると、すんなり理解できるはずです。
また、今回の記事では、行列QとKとVは1組しか登場しませんでしたが、トランスフォーマーでは8組のQKVを使います。これをマルチヘッド・アテンション(多頭アテンション、Multi-head Attention)と呼びます。これによって、多方面からの関連性を知ることができ、良い翻訳が可能となります。
次回は、マルチヘッド・アテンションなどを含めたエンコーダ・ブロックの詳細を解説します。
(続く)
