こんにちは。秋田県のIT企業、北日本コンピューターサービスのR&Dチーム「AUL（アウル）」に所属しています。トラフクロウです。
僕は現在、生成AIの一種であるTransformer（トランスフォーマー）の勉強をしています。いろいろな文献やネット記事を見たことでトランスフォーマーの全体像が少しずつ見えてきました。
（トランスフォーマーの勉強当初に感じた疑問をこちらにまとめています。興味のある方は覗いてみてください。）
トランスフォーマーに関する文献を読んでいると Self-Attention（セルフアテンション）機構 がいつも強調されています。入力単語の中で重要なものに焦点を絞り、文章の意味を特徴づけるための構造です。
計算方法自体はいろいろな所で紹介されているので、意気揚々と読んではみたものの「なんじゃこりゃ？」とつまずいてしまいました。というのも 数式の解読はできるのですが、その式が何をしてくれているのかさっぱりわかりません。
「これはこういうモノなんです。ガーガー言わずに受け入れなさい！」と言われればそれまでですが、ただの計算式と割り切って終わってしまうのは少しさみしい気がします。
製作者の真意はさておきで、セルフアテンションの計算を直感的にイメージできないものか？ 特別な意味を持つ見かたはできないだろうか？ と思うようになりました。
そこで、セルフアテンションの計算では何が起きているのか？ を僕なりにマジメに考えてみることにしました。マジメに考えた結果、少しはセルフアテンションと和解できた気がします。今回の記事ではそれらの内容をまとめていきたいと思います。
セルフアテンションの計算は多角形の収縮と解釈できる
「もったいぶったあげく、フタを開けてみたらよく知られた事実だった」ということになっては目も当てられないので結論だけは先に書いておきます。
【トラフクロウの主張】
トランスフォーマーに入力された文章を高次元空間内の多角形だと考えれば、セルフアテンションの計算は多角形を小さく縮小させる操作であると解釈できる。
以上が今回の記事で一番書きたいことです。賛同が得られるかはわかりませんが、文章を図形だと思えるというのは不思議でワクワクしませんか？。
この記事ではセルフアテンションの計算方法を図形的なイメージで追いかけます。
今回の記事のテーマは次の２つです。
・セルフアテンションの計算式を理解する
・セルフアテンションの計算で何が起きているのか幾何的に眺める
書きたいことを全部書くと結構な量になってしまいそうなので、２回に分けて記事を作りたいと思います。
今回の第１回目で「セルフアテンションの計算方法の解説」を、第２回目では「僕が上の結論に至った経緯」を紹介します。以下お品書きです。
第１回目
・セルフアテンションを味わうための数学的な道具立て
・セルフアテンションの計算方法の解説
第２回目
・トランスフォーマで使われる各種アテンションの計算方法の解説
・セルフアテンションは多角形の収縮で表現できる
今回はセルフアテンションのみにフォーカスします。
トランスフォーマーでは「位置埋め込み」や「残差結合」、「フィードフォワード層」など、たくさんの操作や計算方法を使います。しかし今回は、それらには触れずにセルフテンションの計算方法のみに焦点を当てた内容となっています。
セルフアテンションを味わうための前準備（数学的な道具立て）
ここではセルフアテンションを直感的に理解するために必要な道具や用語を用意します。「知ってるからいいや」という方はスキップしても問題ありません。
「ベクトル」とは矢印であり、数字の並びであり、点である
そんなの知ってるよって人はスキップしてください
最初に用意するものは ベクトル です。「ベクトル」という言葉で何をイメージするかは分野によって流派があります。
ベクトルという言葉が持つ一般的なイメージは 方向 だと思います。「考え方のベクトルが違う」などのベクトルです。基本的な イメージは矢印 になります。
この「方向」に 大きさ という概念を加えて考えるのが物理学を専門とする人たちです。
例えば人が壁を押す動作は矢印（ベクトル）で表現されます。ですが、壁の押すにしてもいろいろな強さで押す場合を考えることができますよね。
物理屋さんはこの力の強さをベクトルの長さで表現し ベクトルの大きさ とよびます。
矢印という見かたから離れて、ベクトルを数字の並びであると解釈する流派もあります。プログラマーや数学を専門とする人たちがそうです。この人たちはいくつかの数字を
$$
[1, 3, -2, 6, -1]
$$
のように並べて「ベクトル（もしくは配列）」とよんでいます。このとき並んだ数字の個数を ベクトルの次元 とよびます。上で挙げたベクトルの次元は5です。
加えて、数学屋さんの場合はさらに別の解釈があります。彼らは数字の並びを見ると 座標 を自然とイメージしてしまいます。$x$軸、$y$軸の座標です。中学校や高校では平面上の点を$(x,y)$で表現するデカルト座標を習います。これも数字の並びなのでプログラマーの流儀では「ベクトル」となります。これが、ベクトルの点としての側面 です。
点としての解釈では「べクトルの次元」は点を打った空間の次元を表します。例えば $(1,2)$ の点は2次元ベクトルなので、2次元平面上の点であることがわかります。
また、平面上の原点 $(0,0)$ を起点に平面上の点に向かって矢印を書くことで数学的なベクトルに、物理学のベクトルの側面が追加されます。このとき「方向」は原点からみて「点」のある方向で、「大きさ」は原点から「点」までの距離になります。
ベクトルの数字を縦に並べるか、横に並べるかにも流派があります。ちなみに僕は縦に並べる派です。今回の記事でも基本的にベクトルは縦に並べる形で書きたいと思います。
記号として縦ベクトルを太字のアルファベットを用いて $\boldsymbol{x}$ のように書き 列ベクトル とよびます。また横ベクトルを表現する場合は $\boldsymbol{x}^T$ と書き 行ベクトル とよびます。
$$
\boldsymbol{x} =
\begin{pmatrix}
1 \
2 \
3 \
\end{pmatrix}, \quad \quad \boldsymbol{x}^T = (1, 2, 3)
$$
「行列」とはベクトルを並べたものである
そんなの知ってるよって人はスキップしてください
数字を並べて作ったものをベクトルとよびましたが、ベクトルそのものを並べることもできます。
列ベクトルを横並びにしたもの、または行ベクトルを縦並びにしたものを 行列 とよびます。ただし、並べるベクトルの次元は同じにする必要があります。
用語として、行列の横方向を 行 、縦方向を 列 とよび、行列内に並んでいる数字１つ１つを 行列の成分 と言います。特に $i$ 行目の $j$ 列目の成分を指定したい場合は $(i,j)$-成分 と表現します。
ベクトル自身も、ベクトルを１つ並べたという意味で行列と言えます。
この記事では、下のように列ベクトルを横に並べて行列とよぶことにします。
$$
A = (\boldsymbol{a}_1, \boldsymbol{a}_2, \boldsymbol{a}_3)
$$
行ベクトルを積み上げたことを強調したいときにはベクトルと同様に $A^T$ と書きます。
$$
A^T =
\begin{pmatrix}
\boldsymbol{a}_1^T \ \
\boldsymbol{a}_2^T \ \
\boldsymbol{a}_3^T
\end{pmatrix}
$$
「ベクトル内積」とはベクトルどうしの相性の良さを表す値である
そんなの知ってるよって人はスキップしてください
$2 \times 2=4$, $2 \times 3=6$ のように数字どうしでは掛け算を考えることができます。一方で、ベクトルどうしでも 内積 とよばれる特別な積（掛け算のようなもの）を考えることができます。
２つの次元が同じベクトル $\boldsymbol{a}, \boldsymbol{b}$ に対してその内積を $\boldsymbol{a} \cdot \boldsymbol{b}$ と書くことにして、計算方法を以下のようにルールとして定めます。
$$
\boldsymbol{a} \cdot \boldsymbol{b} = |\boldsymbol{a}||\boldsymbol{b}| \cos \theta
$$
ここで、$|\boldsymbol{a}|, |\boldsymbol{b}|$ はベクトル $\boldsymbol{a}, \boldsymbol{b}$ それぞれの大きさで、$\theta$ （シータ）は $\boldsymbol{a}, \boldsymbol{b}$ の始点を合わせて絵に書いたときのベクトル間の角度です。この角度 $\theta$ を $\boldsymbol{a}, \boldsymbol{b}$ のなす角 とよびます。
$\cos \theta$ （コサインシータ）は 三角比 とよばれる値です。下の図のよう直角三角形を考えたときに、斜めの辺（斜辺）と底辺の長さの比（分数）を $\cos \theta$ と書くと決められています（斜辺と高さの比をとったものが $\sin \theta$ （サインシータ）です）。
先ほどのベクトル $\boldsymbol{a}, \boldsymbol{b}$ を $\boldsymbol{a}$ が地面と水平になるように回転し、$\boldsymbol{b}$ を斜辺とした直角三角形を考えてみましょう。
このときに、$|\boldsymbol{b}| \cos \theta$ を考えてみると
$$
|\boldsymbol{b}| \cos \theta = |\boldsymbol{b}| \displaystyle{\frac{w}{|\boldsymbol{b}|}} = w
$$
のように計算されます。この結果はベクトル $\boldsymbol{b}$ の終点をベクトル $\boldsymbol{a}$ の上に投影した位置がちょうど長さ $w$ の位置であることを表しています。
このことを意識して $\boldsymbol{a}, \boldsymbol{b}$ の内積を計算すると
$$
\boldsymbol{a} \cdot \boldsymbol{b} = |\boldsymbol{a}||\boldsymbol{b}| \cos \theta = |\boldsymbol{a}||\boldsymbol{b}| \displaystyle{\frac{w}{|\boldsymbol{b}|}} = |\boldsymbol{a}|w
$$
となります。これは、ベクトル $\boldsymbol{a}, \boldsymbol{b}$ の内積 $\boldsymbol{a} \cdot \boldsymbol{b}$ は $\boldsymbol{b}$ の終点を $\boldsymbol{a}$ 上に投影し、そこからさらに $\boldsymbol{a}$ の大きさ倍した位置 になっていることを意味しています。
上の図から、内積の視覚的なイメージがベクトルの投影であることがわかりました。では実際に得られた 内積の値そのものにはどんな意味があるのでしょうか？
分野によってさまざまですが、１つの考え方は ベクトルどうしの仲の良さ です。
内積はベクトルどうしが同じ方向を向いているほど、値が大きくなるという性質を持っています。逆に、違う方向を向いている場合は値が小さくなります（文字どおり方向性の違いで仲が悪いのです）。
また、方向以外の注目ポイントとして、ベクトルの大きさが大きいほど内積の値を大きくすることに貢献できることが挙げられます。ですので、仲の良さに加えて、互いの貢献度 という解釈でもよいかもしれません。
つまるところ、内積を大きくするベクトルどうしは互いに貢献しあい、相性が良いだろう と考えられるのです。
上で見たとおりに内積を計算するには、ベクトル間の角度 $\theta$ を求める必要があります。しかしこれは、とても手間のかかる作業です。
この問題を解決できる素晴らしい性質が知られています。その性質とは、ベクトルの内積は各ベクトルの成分どうしの積の和に等しい というものです。つまり $\boldsymbol{a}^T = (a_1, a_2, a_3), \boldsymbol{b}^T = (b_1, b_2, b_3)$ の内積は
$$
\boldsymbol{a} \cdot \boldsymbol{b} = a_1b_1 + a_2b_2 + a_3b_3
$$
で計算できるということです。このシンプルな計算結果が、不思議なことに、最初に紹介した $\cos \theta$ をつかう複雑な式と等しくなってしまうのです。
「行列の積」はベクトル内積のお化けである
そんなの知ってるよって人はスキップしてください
内積はベクトルの積のようなものでした。行列でも特別な条件下で積を考えることができます。
２つの行列 $A, B$ があったとします。そして $A$ の 行ベクトルの次元 と $B$ の 列ベクトルの次元 が同じとき、$A, B$ の掛け算を次のようにルールとして定めます。
【行列積のルール】
$A$ の $i$ 番目の行ベクトルと $B$ の $j$ 番目の列ベクトルの成分どうしの積の和を $AB$ の $(i, j)$-成分とする。
さて、行列積の計算ルールや数式を観察すると面白い特徴が見えてきます。それは、行列の積は $A$ の行ベクトルと $B$ の列ベクトルの内積を並べて作ることができる ということです。
上の例の行列 $A$ は行ベクトルを縦に積んで
$$
A =
\begin{pmatrix}
\boldsymbol{a}_1^T \ \
\boldsymbol{a}_2^T
\end{pmatrix}
$$
であり、行列 $B$ は列ベクトルを横に並べて
$$
B = (\boldsymbol{b}_1, \boldsymbol{b}_2)
$$
であるとします。すると行列積 $AB$ は
$$
AB =
\begin{pmatrix}
\boldsymbol{a}_1 \cdot \boldsymbol{b}_1 & \boldsymbol{a}_1 \cdot \boldsymbol{b}_2 \
\boldsymbol{a}_2 \cdot \boldsymbol{b}_1 & \boldsymbol{a}_2 \cdot \boldsymbol{b}_2
\end{pmatrix}
$$
と考えることができます。
ベクトル内積は「ベクトルどうしの相性の良さ」と考えることができるので、行列積は２つの行列 $A, B$ 構成するすべてのベクトルどうしの相性の良さを並べたもの であると解釈することができます。
行列 $A$ が $k$ 本の行と $l$ 本の列を持っているとき、行列 $A$ は $(k,l)$-型の行列である とよばれます。行列積のルールから $A$ に掛けられる行列 $B$ は $(l,n)$-型である必要があります。そして積を計算して得られる新たな行列 $AB$ は $(k,n)$-型になることが知られています。
この事実を知ったうえで内積を振り返ると、内積は $(1,l)$-型の行列（行ベクトル）と $(l,1)$-型の行列（列ベクトル）の行列積と考えることもできます。よくできているなと感心しますよね。
「線形変換」はベクトルをグリグリ動かす操作である
そんなの知ってるよって人はスキップしてください
$(k, l)$-型の行列と $(l, 1)$-型のベクトルの間で行列の掛け算を行うことで、新しい $(k, 1)$-型のベクトルを手に入れることができます。視覚的なイメージとしては、ベクトルに行列を掛け算することで別のベクトルにうつすことができるわけです。
行列を掛ける前後でベクトルの次元が変化しないとき、この掛け算を 線形変換 とよびます。$n$ 次元ベクトルを $n$ 次元空間内の点であると考えると $(n, n)$-型の行列による線形変換は空間内で点を移動させることと同じです。
次の図は平面上の点（ベクトル）の集まりに特定の行列を掛け算した結果の図です。
この結果から、線形変換されたベクトルは同じ空間の中で伸び縮みや回転をすることが分かります。また違うベクトルどうしに同じ行列を掛けると他のベクトルも同じように変換されることも分かりますね。
つまり、線形変換をすることでベクトルたちが同じ規則で空間内を動く ということです。
「ソフトマックス関数」は複数の数値を確率っぽく変換する便利なヤツ
そんなの知ってるよって人はスキップしてください
名前に「関数」と付くだけにソフトマックス関数は特別な関数の名前です。「関数」という言葉が何を意味するかは分野によって流派があります。ここでの関数は、入力を受けとりガリガリ計算して特定の形で出力するモノという認識で問題ありません。
ソフトマックス関数はベクトルを受け取り、同じ次元のベクトルを出力します。このときの入力ベクトルを
$$
\boldsymbol{x}^T = (x_1, x_2, x_3, x_4)
$$
として、出力ベクトルを
$$
\boldsymbol{y}^T = (y_1, y_2, y_3, y_4)
$$
とします。各 $y_i$ を計算する具体的な規則は次のように決まっています。
$$
\begin{align*}
y_i &= \displaystyle{\frac{e^{x_i}}{\sum_{i=1}^4 e^{x_i}}} \
&= \displaystyle{\frac{e^{x_i}}{e^{x_1} + e^{x_2} + e^{x_3} + e^{x_4}}} \quad \quad (i = 1, 2, 3, 4)
\end{align*}
$$
すこし式がややこしいので順番に見ていきましょう。
まず、$\sum$ は「シグマ」と読みます。$\sum$ は、右側に書いてある式を何回か足し合わせることを意味する記号です。上では $e^{x_i}$ を足しています。
$e^{x_i}$を何回足せばよいかは $\sum$ の下と上の数字を見るとわかります。下が開始の番号で上が終了の番号です。そのため上の式は、「$e^{x_i}$ を $i=1$ から $i=4$ まで足しわせなさい」という意味になります（入力が $n$ 次元ベクトルの場合は $n$ 個の足し合わせになります）。
ここで出てくる $e$ とは ネイピア数 とよばれている数字です。具体的には
$$
e = 2.718281828459 \cdots
$$
と無限に続く数です。深入りはしませんが、ネイピア数を使うことで微積分などを行う際にいくらかの恩恵を得ることができます。
$e^{x_i}$ はネイピア数 $e$ を $x_i$ 乗することを意味しています。このように $e$ を何かしらの数でべき乗した $e^x$ を指数関数とよびます。下は指数関数のグラフです。
グラフの見かたは次のとおりです。まず横軸 $x$ は入力値をずらっと並べたもの（並べすぎて線に見えている）。グラフ自身は、特定の $x$ で $e^x$ を計算した結果（これも点を並べすぎて線に見えている）です。
こうしてみると 指数関数は、どんな $x$ を入れても必ず $0$ より大きい数に変換できる ことがわかります。そして、入力値である $x$ の大小関係を変換後に変化させない 性質を持ちます。
つまり入力前に $x_1 < x_2$ だった入力値は変換後も大小関係を保ち $e^{x_1} < e^{x_2}$ となるということです。
さて、これでソフトマックス関数の $y_i$ が読めるようになりました。
$$
\begin{align*}
y_i &= \displaystyle{\frac{e^{x_i}}{\sum_{i=1}^4 e^{x_i}}} \
&= \displaystyle{\frac{e^{x_i}}{e^{x_1} + e^{x_2} + e^{x_3} + e^{x_4}}} \quad \quad (i = 1, 2, 3, 4)
\end{align*}
$$
まとめると、ソフトマックス関数は入力された全ての $x_i$ を指数関数に通し、特定の $e^{x_i}$ が全体の内どの程度の割合を占めるかを考えている わけです。さらに入力値の大小関係を保存できるという特性から 大きな $x_i$ をより際立たせ、小さな $x_i$ の影響をより目立たなくする という特徴があります。
計算式から明らかですが、全ての $y_i$ の総和は $1$ になります。この性質は確率が持っている性質の１つです。このことから ソフトマックス関数は入力値に対して疑似的な確率を割り振る道具 としても使われることがあります。
セルフアテンションの計算方法の紹介
トランスフォーマーに入力された文章は意味の最小単単位である単語に分割され、各単語には適当なベクトルが割り当てられます（詳しくはこちら）。
各単語ベクトルは互いの関係性を評価されながら、より洗練されたベクトルにつくり変えられます。このとき計算に使われる 関係性の強さを表す数値がセルフアテンション です。
例えば、上の図では「これ」という単語のベクトルを作るために、他の単語と「これ」の関係性が数値化されています。そして、この関係性をつなぎとして複数の単語ベクトルを１つに束ねることを考えます。数式にすると下のようなイメージです。
\boldsymbol{x}^{\prime}_{これ} = 0.6 \boldsymbol{x}_{これ} + 0.2 \boldsymbol{x}_{は} + 0.8 \boldsymbol{x}_{ペン} + 0.1 \boldsymbol{x}_{です}
ここで作られた新しい $\boldsymbol{x}_{これ}^{\prime}$ が文章全体を考慮して作られた「これ」ベクトルになります。つまりは、より洗練された「これ」ベクトルです。
それでは、以上の流れを具体的な計算式で見ていきましょう。
キー、クエリ、バリュー
まずは、入力された各単語ベクトル $\boldsymbol{x}_i$ に対し キー、クエリ、バリュー とよばれる３種類のベクトルが作られます（$\boldsymbol{x}_i$ は既に位置埋め込みベクトルが反映されているものとします）。
簡単のため $\boldsymbol{x}_i$ を $2$ 次元ベクトルであるとします。このベクトルに３種類の $(2, 2)$-型行列 $W_K, W_Q, W_V$ をそれぞれ掛け算して３本のベクトルを作ります。それらベクトルを $\boldsymbol{k}_i, \boldsymbol{q}_i, \boldsymbol{v}_i$ と書くことにしましょう。
$$
\boldsymbol{k}_i = W_K \boldsymbol{x}_i, \quad \quad \boldsymbol{q}_i = W_Q \boldsymbol{x}_i, \quad \quad \boldsymbol{v}_i = W_V \boldsymbol{x}_i
$$
この $\boldsymbol{k}_i, \boldsymbol{q}_i, \boldsymbol{v}_i$ がそれぞれ、キー、クエリ、バリューベクトルです。
ここで、これらの行列によるベクトルの計算は 線形変換 です。そのため、このキー、クエリ、バリューベクトルを求める操作は、単語ベクトル $\boldsymbol{x}_i$ を異なる観点で変換し個別の特徴を待たせている のだと予想することができます。
今後の計算のために各単語に対応するキー、クエリ、バリューベクトルをそれぞれ種類ごとにまとめて行列で書いておきます。「これはペンです」の例文に習ってベクトルは4本であるとしましょう。
$$
K = (\boldsymbol{k}_1, \boldsymbol{k}_2, \boldsymbol{k}_3, \boldsymbol{k}_4), \quad Q = (\boldsymbol{q}_1, \boldsymbol{q}_2, \boldsymbol{q}_3, \boldsymbol{q}_4) , \quad V = (\boldsymbol{v}_1, \boldsymbol{v}_2, \boldsymbol{v}_3, \boldsymbol{v}_4)
$$
重要度（アテンション）の計算と単語ベクトルの生成
最終的に得られる、洗練された単語ベクトルを並べた行列を $A$ と書くことにします。すると、行列 $A$ は キー、クエリ、バリュー行列 $K, Q, V$ を使って次のように計算されます。
$$
A = V \cdot \text{softmax} \Biggl( \displaystyle{\frac{K^T Q}{\sqrt{D}}} \Biggr) \quad \quad （ D は単語ベクトルの次元、今の例では 2 ）
$$
元論文との表記の違いは流派の問題です。
上の式では、行列を列ベクトルを並べたものとしているため、元論文の「Attention Is All You Need」に掲載されている式とは表記のずれがあります。論文の著者は行ベクトルを積み上げる派だったみたいですね・・・・・・
僕はこの式を見ても「結局何が起きてるの？」といった印象でした。ここからは数式の持つイメージを意識しながらこの式を順番に読み解いていこうと思います。
まず、「softmax」 と書いてある部分はソフトマックス関数です。ソフトマックス関数に何やら複雑なモノが入力されていますね。
ソフトマックス関数への入力部分をこれまでの表記で書き下すと次のようになります。
\begin{align*}
\displaystyle{\frac{K^T Q}{\sqrt{2}}} &= \displaystyle{\frac{1}{\sqrt{2}}} K^T Q \\
& = \displaystyle{\frac{1}{\sqrt{2}}}
\begin{pmatrix}
\boldsymbol{k}_1^T \\ \\
\boldsymbol{k}_2^T \\ \\
\boldsymbol{k}_3^T \\ \\
\boldsymbol{k}_3^T
\end{pmatrix} (\boldsymbol{q}_1, \boldsymbol{q}_2, \boldsymbol{q}_3, \boldsymbol{q}_4) \\
&= \displaystyle{\frac{1}{\sqrt{2}}}
\begin{pmatrix}
\boldsymbol{k}_1 \cdot \boldsymbol{q}_1 & \boldsymbol{k}_1 \cdot \boldsymbol{q}_2 & \boldsymbol{k}_1 \cdot \boldsymbol{q}_3 & \boldsymbol{k}_1 \cdot \boldsymbol{q}_4 \\
\boldsymbol{k}_2 \cdot \boldsymbol{q}_1 & \boldsymbol{k}_2 \cdot \boldsymbol{q}_2 & \boldsymbol{k}_2 \cdot \boldsymbol{q}_3 & \boldsymbol{k}_2 \cdot \boldsymbol{q}_4 \\
\boldsymbol{k}_3 \cdot \boldsymbol{q}_1 & \boldsymbol{k}_3 \cdot \boldsymbol{q}_2 & \boldsymbol{k}_3 \cdot \boldsymbol{q}_3 & \boldsymbol{k}_3 \cdot \boldsymbol{q}_4 \\
\boldsymbol{k}_4 \cdot \boldsymbol{q}_1 & \boldsymbol{k}_4 \cdot \boldsymbol{q}_2 & \boldsymbol{k}_4 \cdot \boldsymbol{q}_3 & \boldsymbol{k}_4 \cdot \boldsymbol{q}_4
\end{pmatrix}
\end{align*}
行列の中身は各キーベクトルとクエリベクトルの内積です。つまり キーベクトルとクエリベクトルの相性の良さを表す数字が並んでいる ことになります。下の図は１列目の内積計算のイメージです。
クエリベクトルに注目して最後の行列を見てみると、行列の１列目に１番目の単語と他の単語の相性が並んでいることがわかります。表記を簡単にするために
\boldsymbol{k}_i \cdot \boldsymbol{q}_j = c_{ij}
とおき、各列ベクトルを $\boldsymbol{c}_i$ と書くことにしましょう。すると
\begin{align*}
\displaystyle{\frac{1}{\sqrt{2}}}
\begin{pmatrix}
\boldsymbol{k}_1 \cdot \boldsymbol{q}_1 & \boldsymbol{k}_1 \cdot \boldsymbol{q}_2 & \boldsymbol{k}_1 \cdot \boldsymbol{q}_3 & \boldsymbol{k}_1 \cdot \boldsymbol{q}_4 \\
\boldsymbol{k}_2 \cdot \boldsymbol{q}_1 & \boldsymbol{k}_2 \cdot \boldsymbol{q}_2 & \boldsymbol{k}_2 \cdot \boldsymbol{q}_3 & \boldsymbol{k}_2 \cdot \boldsymbol{q}_4 \\
\boldsymbol{k}_3 \cdot \boldsymbol{q}_1 & \boldsymbol{k}_3 \cdot \boldsymbol{q}_2 & \boldsymbol{k}_3 \cdot \boldsymbol{q}_3 & \boldsymbol{k}_3 \cdot \boldsymbol{q}_4 \\
\boldsymbol{k}_4 \cdot \boldsymbol{q}_1 & \boldsymbol{k}_4 \cdot \boldsymbol{q}_2 & \boldsymbol{k}_4 \cdot \boldsymbol{q}_3 & \boldsymbol{k}_4 \cdot \boldsymbol{q}_4
\end{pmatrix} &= \displaystyle{\frac{1}{\sqrt{2}}} (\boldsymbol{c}_1, \boldsymbol{c}_2, \boldsymbol{c}_3, \boldsymbol{c}_4) \\
&= \biggl(\displaystyle{\frac{1}{\sqrt{2}}} \boldsymbol{c}_1, \displaystyle{\frac{1}{\sqrt{2}}} \boldsymbol{c}_2, \displaystyle{\frac{1}{\sqrt{2}}} \boldsymbol{c}_3, \displaystyle{\frac{1}{\sqrt{2}}} \boldsymbol{c}_4 \biggr)
\end{align*}
となります。したがってソフトマックス関数部分の全体像は
$$
\text{softmax} \Biggl( \displaystyle{\frac{K^T Q}{\sqrt{D}}} \Biggr) =
\text{softmax} \biggl(\displaystyle{\frac{1}{\sqrt{2}}} \boldsymbol{c}_1, \displaystyle{\frac{1}{\sqrt{2}}} \boldsymbol{c}_2, \displaystyle{\frac{1}{\sqrt{2}}} \boldsymbol{c}_3, \displaystyle{\frac{1}{\sqrt{2}}} \boldsymbol{c}_4 \biggr)
$$
となっていたわけです。
さて、ソフトマックス関数はベクトル受け取りベクトルの成分の総和が $1$ になるベクトルを出力します。今回は列ベクトルごとにソフトマックス関数を適用します。すなわち
\begin{align*}
& \text{softmax} \biggl(\displaystyle{\frac{1}{\sqrt{2}}} \boldsymbol{c}_1, \displaystyle{\frac{1}{\sqrt{2}}} \boldsymbol{c}_2, \displaystyle{\frac{1}{\sqrt{2}}} \boldsymbol{c}_3, \displaystyle{\frac{1}{\sqrt{2}}} \boldsymbol{c}_4 \biggr) \\
&= \biggl( \text{softmax} \biggl( \displaystyle{\frac{1}{\sqrt{2}}} \boldsymbol{c}_1 \biggr), \text{softmax} \biggl( \displaystyle{\frac{1}{\sqrt{2}}} \boldsymbol{c}_2 \biggr), \text{softmax} \biggl( \displaystyle{\frac{1}{\sqrt{2}}} \boldsymbol{c}_3 \biggr), \text{softmax} \biggl( \displaystyle{\frac{1}{\sqrt{2}}} \boldsymbol{c}_4 \biggr) \biggr)
\end{align*}
です。$\boldsymbol{c}_i$ はi番目の単語と他の単語の相性を表すベクトルでした。これをソフトマックス関数に通すことで
\text{softmax} \biggl( \displaystyle{\frac{1}{\sqrt{2}}} \boldsymbol{c}_i^T \biggr) =
\begin{pmatrix}
c_{11}^{\prime} \\
c_{12}^{\prime} \\
c_{13}^{\prime} \\
c_{14}^{\prime}
\end{pmatrix}
= \boldsymbol{c^{\prime}}_i
が得られます。ソフトマックス関数は入力値の大小を顕著にする性質があったので、$\boldsymbol{c^{\prime}}_i$ は $i$ 番目の単語と相性の良かった部分をより強調したベクトルとなっているわけです。
ここで、相性ベクトル $\boldsymbol{c}_i$ が次元のルートで割られている理由を補足しておきます。上でも書きましたが、ソフトマックス関数は入力された大きな値をより大きく際立たせます。
一方でベクトル内積の事情として、ベクトルの次元が高くなると内積の値も大きくなりやすいという傾向があります。これは内積の計算がベクトルの次元数に応じて増えることからも想像しやすいです。
そのため、ベクトルの次元が大きくなったとき、単語の組み合わせによっては内積の値に大きな開きができてしまい、ソフトマックス関数によってその格差がさらに広がってしまいます。
これは偏った学習の原因になります。これを防ぐために、内積を次元のルートで割り、ソフトマックス関数への入力値の格差が極端に大きくならないようにしているのです。
ここまでの内容をまとめると、元々の式が
$$
A = V \text{softmax} \Biggl( \displaystyle{\frac{K^T Q}{\sqrt{D}}} \Biggr) =
V (\boldsymbol{c^{\prime}}_1, \boldsymbol{c^{\prime}}_2, \boldsymbol{c^{\prime}}_3, \boldsymbol{c^{\prime}}_4)
$$
と書けることがわかりました。あとは素直に行列の掛け算を行うだけです。
行列積のルールに従って、バリューベクトル $V$ の行と、$\boldsymbol{c^{\prime}}_i$ の内積を並べていきます。計算結果を列ベクトルの形でまとめると次のようになります。
A = \biggl(\sum_{l=1}^4 c_{1l}^{\prime} \boldsymbol{v}_l, \sum_{l=1}^4 c_{2l}^{\prime} \boldsymbol{v}_l, \sum_{l=1}^4 c_{3l}^{\prime} \boldsymbol{v}_l, \sum_{l=1}^4 c_{4l}^{\prime} \boldsymbol{v}_l \biggr)
列ベクトルが和の形でかかれているので少し見ずらいですね。最初の成分だけ取り出してみると次の通りです。
\sum_{l=1}^4 c_{1l}^{\prime} \boldsymbol{v}_l = c_{11}^{\prime} \boldsymbol{v}_1 + c_{12}^{\prime} \boldsymbol{v}_2 + c_{13}^{\prime} \boldsymbol{v}_3 + c_{14}^{\prime} \boldsymbol{v}_4
$c_{1l}^{\prime}$ は1番目の単語と $l$ 番目の単語の相性でした。また各 $\boldsymbol{v}_l$ は $l$ 番目の単語ベクトルの代用品（線形変換した単語ベクトル）です。
したがって、上の足し算で得られるベクトルは各単語ベクトルを単語どうしの相性（関係性）を考慮して足し合わせたベクトルになっています。
これが、冒頭で述べた「洗練された単語ベクトル」です。以上がトランスフォーマーで採用されているセルフアテンションの計算になります。
単語の 重要度に注目して計算をするという気持ちがなんとなく見えた のでないでしょうか。
第１回目はここまでです。最後まで読んでいただきありがとうございました！
前回に引き続き、またしても膨大な量になってしまいました。数式がメインの説明だったので読みづらい部分も多かったと思いますが、最後まで読んでいただきありがとうございました。
セルフアテンションを調べる中で、数式が簡単な線形代数の式のみで記述されていることに気が付きました。そこで、数式の持つ幾何学的なイメージを持ち出せばもっと直感的にトランスフォーマーを理解できるのでは？と思ったことが今回の記事の始まりです。ですがやはり一筋縄ではいきませんね。
もっとわかりやすい見かたや、面白い考え方があれば教えていただけるとうれしいです。
次回は、今回紹介しきれなかったマルチヘッドアテンションや、記事の冒頭で述べた多角形を使ったセルフアテンションのイメージについて書きたいと思います。
参考文献ならびに参考図書
- Attention Is All You Need
- 山田育矢, 鈴木正敏, 山田康輔, 李凌寒, 「大規模言語モデル入門」, 技術評論社, 2023