自己注意：スケール付き内積注意・マルチヘッド・因果マスク

0. この記事の要点

文章を読むとき、私たちは「それ」や「かける」の意味を、文のほかの場所を見て決めています。自己注意は、計算機が同じことをするための部品です。列の中の各語について、どの語をどれだけ参考にするかを数で決め、その割合で情報を混ぜ合わせます。この記事では、この演算がどんな範囲の値を出すか、語順をどう扱うか、未来の語を見ないようにするにはどうするかを、式で確かめます。

• 自己注意は、列の各位置のベクトルを、ほかの位置のベクトルの重みつき平均で置き換える演算です。重みは内積を softmax に通して作るので、非負で和が 1 です。出力の各行は値ベクトルの凸結合（重みが非負で和が 1 の重みつき平均）になります（命題 4.2）。
• softmax は平行移動で変わらず、大小を保ち、置換と交換します（命題 3.1）。ヤコビ行列（入力を少し動かしたとき出力がどう動くかを表す行列）は $\mathrm{diag}(\boldsymbol{\sigma}) - \boldsymbol{\sigma}\boldsymbol{\sigma}^{\mathsf{T}}$ で、出力が 1 点に集中するほど小さくなります（命題 3.2）。内積を $\sqrt{d_k}$ で割るのは、この集中を次元に依らず抑えるためです（命題 4.4）。
• 位置情報もマスクもない自己注意は、入力の行を並べ替えると出力の行が同じように並べ替わるだけです（定理 5.2）。語順を区別できないので、前章の位置符号化が必要になります。
• 因果マスクを入れると、出力の第 $i$ 行は $\boldsymbol{x}_1, \ldots, \boldsymbol{x}_i$ だけの関数になります（定理 6.3）。これにより、長さ $n$ の列に対する 1 回の計算で、$n$ 個の「接頭辞から次のトークン」の予測が同時に得られます（系 6.4）。
• 1 層の乗算の回数は $4nd^2 + 2n^2 d$、パラメータ数は $4d^2 + 4d$ で、どちらもヘッド数に依りません（命題 7.2、命題 7.1）。

この記事で使う事前知識

• 行列の積: $p \times q$ 行列と $q \times r$ 行列から $p \times r$ 行列を作る演算で、成分は行と列の対応する成分の積の和。→ 行列の積
• 期待値: 確率変数の値を確率で重みづけて平均したもの。→ 期待値
• 分散: 平均からのずれの 2 乗の期待値で、値の散らばりの大きさを表す。→ 分散
• 指数関数: $e^{x}$ のこと。常に正で、$x$ が増えると増える。→ 指数関数
• 位置符号化: トークンの並び順の情報をベクトルに持たせる仕組み。前章で扱いました。→ 位置符号化

1. 動機：位置どうしで情報をやり取りする

前章までで、トークン列は $n \times d$ の行列 $X$ になりました。第 $i$ 行 $\boldsymbol{x}_i$ は第 $i$ トークンを表すベクトルです。この時点では、各行はそのトークン 1 個の情報しか持っていません。「それ」が何を指すか、「かける」が電話なのか橋なのかを決めるには、ほかの位置の情報を取り込む必要があります。

取り込み方として、どの位置からどれだけ取り込むかを入力に応じて決めるのが注意です。位置 $i$ が「こういう情報が欲しい」というベクトル（クエリ）を出し、各位置 $j$ が「自分はこういう情報を持っている」というベクトル（キー）と、渡す中身（バリュー）を出します。クエリとキーの内積が大きい位置ほど大きな重みを付けて、バリューを平均します。注意の仕組みは Bahdanau らが機械翻訳で導入し（[2]）、Vaswani らは再帰も畳み込みも使わず注意だけでモデルを組みました（[1]）。

この章では、定義のあとで次の 3 つの問いに命題として答えます。

1. 注意の出力はどんな範囲に収まるか（凸結合）。
2. 入力の順番を入れ替えると出力はどうなるか（同変性）。これが前章の位置符号化が必要な理由です。
3. 「未来を見ない」という制約はどう入れ、それで何が保証されるか（因果マスク）。これが言語モデルの学習を並列にできる理由で、第 6 章の KV キャッシュの正しさの根拠にもなります。

2. 準備：記号・置換行列・softmax

行列 $X \in \mathbb{R}^{n \times d}$ の第 $i$ 行を列ベクトルとみなしたものを $\boldsymbol{x}_i \in \mathbb{R}^{d}$ と書きます。全成分が 1 のベクトルを $\boldsymbol{1}$ と書きます。行列の積（定義 4.1［行列と連立一次方程式］）について、$XW$ の第 $i$ 行は $W^{\mathsf{T}}\boldsymbol{x}_i$ です。

定義 2.1（置換行列・同変・不変）↗

$\{1, \ldots, n\}$ の置換 $\pi$（定義 2.1［行列式とその性質］）に対し、$(P_\pi)_{ij} := 1$（$j = \pi(i)$ のとき）、$0$（それ以外）で定まる $n$ 次の正方行列 $P_\pi$ を置換行列という。

写像 $f : \mathbb{R}^{n \times d} \to \mathbb{R}^{n \times d'}$ が、すべての置換行列 $P$ とすべての $X$ について $f(PX) = P f(X)$ を満たすとき置換同変、$f(PX) = f(X)$ を満たすとき置換不変という。

関数 $g : \mathbb{R}^{d} \to \mathbb{R}^{d'}$ に対し、各行に $g$ を施す写像 $G(X)_i := g(\boldsymbol{x}_i)$ を行ごとに働く写像という（$G(X)_i$ は $G(X)$ の第 $i$ 行）。

イメージをつかむために小さい例を置きます。$n = 3$ で、$\pi(1) = 2$、$\pi(2) = 1$、$\pi(3) = 3$（1 番目と 2 番目の入れ替え）とすると、$P_\pi X$ の行は上から順に $X$ の第 2 行、第 1 行、第 3 行になります。置換行列は「行を並べ替える操作」を行列で書いたものです。

置換行列について、定義から直接出る事実を 3 つ確認しておきます。

• $(P_\pi X)_{ij} = \sum_l (P_\pi)_{il} X_{lj} = X_{\pi(i) j}$。つまり $P_\pi X$ の第 $i$ 行は $X$ の第 $\pi(i)$ 行で、$P_\pi$ を左から掛けると行が並べ替わります。ベクトルなら $(P_\pi \boldsymbol{z})_i = z_{\pi(i)}$ です。
• $(A P_\pi^{\mathsf{T}})_{ij} = \sum_l A_{il}(P_\pi)_{jl} = A_{i\pi(j)}$。右から転置を掛けると列が並べ替わります。
• $(P_\pi^{\mathsf{T}} P_\pi)_{jl} = \sum_i (P_\pi)_{ij}(P_\pi)_{il}$ で、$(P_\pi)_{ij}(P_\pi)_{il} = 1$ となるのは $j = \pi(i) = l$ のときだけです。$\pi$ は全単射なので、$j = l$ ならそのような $i$ がちょうど 1 つあり、$j \ne l$ なら 1 つもありません。よって $P_\pi^{\mathsf{T}} P_\pi = I$。また各行に 1 がちょうど 1 個あるので $P_\pi \boldsymbol{1} = \boldsymbol{1}$ です。

定義 2.2（softmax とマスクつき softmax）↗

$\boldsymbol{z} \in \mathbb{R}^{n}$ に対し

$$\sigma(\boldsymbol{z})_i := \frac{e^{z_i}}{\sum_{j=1}^{n} e^{z_j}} \qquad (i = 1, \ldots, n)$$

で定まる写像 $\sigma : \mathbb{R}^n \to \mathbb{R}^n$ を softmax 関数という。行列に対しては行ごとに施す。

空でない部分集合 $U \subset \{1, \ldots, n\}$ に対し、マスクつき softmax $\sigma^{U} : \mathbb{R}^n \to \mathbb{R}^n$ を

$$\sigma^{U}(\boldsymbol{z})_i := \frac{e^{z_i}}{\sum_{j \in U} e^{z_j}} \quad (i \in U), \qquad \sigma^{U}(\boldsymbol{z})_i := 0 \quad (i \notin U)$$

で定める。$U$ に入らない添字を「マスクされた」成分という。「マスクされた成分に $-\infty$ を足してから softmax を取る」という書き方は、$\sigma^{U}$ のことを指すものとする。

$-\infty$ は実数ではないので、「足す」という書き方はそのままでは定義になりません。上では $\sigma^U$ を直接定義し、それが極限として正当化できることを 命題 3.3 で示します。

3. softmax の性質

命題 3.1（softmax の基本性質）✓

$\boldsymbol{z} \in \mathbb{R}^n$ とする。

1. すべての $i$ で $\sigma(\boldsymbol{z})_i > 0$ であり、$\sum_{i=1}^{n} \sigma(\boldsymbol{z})_i = 1$。
2. （平行移動不変）任意の $c \in \mathbb{R}$ について $\sigma(\boldsymbol{z} + c\boldsymbol{1}) = \sigma(\boldsymbol{z})$。
3. （大小を保つ）$z_i < z_j$ ならば $\sigma(\boldsymbol{z})_i < \sigma(\boldsymbol{z})_j$、$z_i = z_j$ ならば $\sigma(\boldsymbol{z})_i = \sigma(\boldsymbol{z})_j$。
4. （置換同変）任意の置換行列 $P$ について $\sigma(P\boldsymbol{z}) = P\sigma(\boldsymbol{z})$。

証明（命題 3.1）

$S := \sum_{j} e^{z_j}$ と置きます。指数関数は正の値を取り、狭義単調増加で、$e^{x+y} = e^x e^y$ を満たします（定理 3.3［指数・対数・三角関数をきちんと定義する］ の 2 と 4）。

1. $e^{z_i} > 0$、$S > 0$ より $\sigma(\boldsymbol{z})_i = e^{z_i}/S > 0$。和は $\sum_i e^{z_i}/S = S/S = 1$。
2. $\sigma(\boldsymbol{z} + c\boldsymbol{1})_i = \dfrac{e^{z_i + c}}{\sum_j e^{z_j + c}} = \dfrac{e^{c}e^{z_i}}{e^{c}\sum_j e^{z_j}} = \sigma(\boldsymbol{z})_i$。
3. $z_i < z_j$ なら $e^{z_i} < e^{z_j}$ で、同じ正の数 $S$ で割っても不等号は変わりません。等しい場合も同様です。
4. $P = P_\pi$ とすると $(P\boldsymbol{z})_i = z_{\pi(i)}$ です。$\pi$ は全単射なので $\sum_j e^{z_{\pi(j)}} = \sum_{j'} e^{z_{j'}} = S$（和の順番を入れ替えただけ）。よって $\sigma(P\boldsymbol{z})_i = e^{z_{\pi(i)}}/S = \sigma(\boldsymbol{z})_{\pi(i)} = (P\sigma(\boldsymbol{z}))_i$。

∎

2 は実装で重要です。$z_i$ が大きいと $e^{z_i}$ は浮動小数点数の範囲を超えますが、最大値を引いてから指数を取れば、すべての指数が 1 以下になり、結果は変わりません。8 節のコードもそうしています。

次の命題のねらいは、softmax に勾配がどれだけ流れるかを調べることです。証明の 1 では、$\boldsymbol{v}^{\mathsf{T}}J\boldsymbol{v}$ が「$\boldsymbol{\sigma}$ を確率とみなしたときの $v_i^2$ の平均から、$v_i$ の平均の 2 乗を引いたもの」、つまり分散の形になることを使い、それが 0 以上であることをコーシー・シュワルツの不等式で示します。

命題 3.2（softmax のヤコビ行列）✓

$\boldsymbol{\sigma} := \sigma(\boldsymbol{z})$ と書く。softmax は $\mathbb{R}^n$ 上で微分可能で、

$$\frac{\partial \sigma_i}{\partial z_j} = \sigma_i(\delta_{ij} - \sigma_j), \qquad \text{すなわち}\quad J(\boldsymbol{z}) = \mathrm{diag}(\boldsymbol{\sigma}) - \boldsymbol{\sigma}\boldsymbol{\sigma}^{\mathsf{T}}$$

である（$\delta_{ij}$ は $i = j$ のとき 1、それ以外 0）。さらに次が成り立つ。

1. $J(\boldsymbol{z})$ は対称で半正定値である。
2. $J(\boldsymbol{z})\boldsymbol{1} = \boldsymbol{0}$。
3. すべての $i, j$ で $|J(\boldsymbol{z})_{ij}| \le 1 - \max_k \sigma_k$。

証明（命題 3.2）

$S = \sum_k e^{z_k}$ とすると $\partial S/\partial z_j = e^{z_j}$、$\partial e^{z_i}/\partial z_j = \delta_{ij}e^{z_i}$ です。商の微分より

$$\frac{\partial \sigma_i}{\partial z_j} = \frac{\delta_{ij}e^{z_i}\, S - e^{z_i} e^{z_j}}{S^2} = \delta_{ij}\sigma_i - \sigma_i\sigma_j = \sigma_i(\delta_{ij} - \sigma_j).$$

行列で書くと、$\delta_{ij}\sigma_i$ が $\mathrm{diag}(\boldsymbol{\sigma})$、$\sigma_i\sigma_j$ が $\boldsymbol{\sigma}\boldsymbol{\sigma}^{\mathsf{T}}$ の $(i, j)$ 成分です。

1. $\delta_{ij}\sigma_i - \sigma_i\sigma_j$ は $i$ と $j$ を入れ替えても同じなので対称です。任意の $\boldsymbol{v} \in \mathbb{R}^n$ について

$$\boldsymbol{v}^{\mathsf{T}} J \boldsymbol{v} = \sum_i \sigma_i v_i^2 - \Big(\sum_i \sigma_i v_i\Big)^2$$

です。$a_i := \sqrt{\sigma_i}$、$b_i := \sqrt{\sigma_i}\,v_i$ にコーシー・シュワルツの不等式（定理 4.2［内積空間とグラム・シュミット直交化］）を使うと $\big(\sum_i \sigma_i v_i\big)^2 = \big(\sum_i a_i b_i\big)^2 \le \big(\sum_i \sigma_i\big)\big(\sum_i \sigma_i v_i^2\big) = \sum_i \sigma_i v_i^2$（命題 3.1 の 1 で $\sum_i\sigma_i = 1$）。よって $\boldsymbol{v}^{\mathsf{T}}J\boldsymbol{v} \ge 0$ で、半正定値です（定義 6.2［スペクトル定理］）。

2. 第 $i$ 成分は $\sum_j \sigma_i(\delta_{ij} - \sigma_j) = \sigma_i - \sigma_i\sum_j\sigma_j = 0$。

3. $\sigma_k$ が最大になる添字の 1 つを $k^{*}$ とします。$i \ne j$ のとき、$i, j$ の少なくとも一方は $k^{*}$ ではありません。$i \ne k^{*}$ としてよく（$J$ は対称）、$|J_{ij}| = \sigma_i\sigma_j \le \sigma_i \le \sum_{k \ne k^{*}}\sigma_k = 1 - \sigma_{k^{*}}$。$i = j$ のとき $J_{ii} = \sigma_i(1 - \sigma_i)$ で、$i = k^{*}$ なら $\sigma_{k^{*}}(1 - \sigma_{k^{*}}) \le 1 - \sigma_{k^{*}}$、$i \ne k^{*}$ なら $\sigma_i(1 - \sigma_i) \le \sigma_i \le 1 - \sigma_{k^{*}}$ です。

∎

2 は平行移動不変性（命題 3.1 の 2）の微分形です。$\boldsymbol{1}$ の向きに動かしても出力が変わらないので、その向きの微分は 0 になります。3 は、softmax の出力がほぼ 1 点に集中している（$\max_k \sigma_k$ が 1 に近い）と、ヤコビ行列の全成分が 0 に近くなることを言っています。そうなると softmax を通る勾配がほとんど流れません。交差エントロピー（予測した確率が正解からどれだけ外れているかを測る損失）と組み合わせたときの勾配 $\boldsymbol{\sigma} - \boldsymbol{y}$（命題 7.1［ニューラルネットワークと逆伝播］）は、このヤコビ行列に損失の勾配を掛けた結果です（ヤコビ行列の式 $\mathrm{diag}(\boldsymbol{\sigma}) - \boldsymbol{\sigma}\boldsymbol{\sigma}^{\mathsf{T}}$ そのものは同じ命題にも述べられています。ここで新しく示したのは 1〜3 の性質です）。

命題 3.3（マスクは極限として正当化できる）✓

$U \subset \{1, \ldots, n\}$ を空でない部分集合とし、$\boldsymbol{m} \in \mathbb{R}^n$ を、$i \notin U$ のとき $m_i = 1$、$i \in U$ のとき $m_i = 0$ であるベクトルとする。任意の $\boldsymbol{z} \in \mathbb{R}^n$ について

$$\lim_{c \to \infty} \sigma(\boldsymbol{z} - c\,\boldsymbol{m}) = \sigma^{U}(\boldsymbol{z})$$

が成り立つ。さらに $i \notin U$ のとき $0 < \sigma(\boldsymbol{z} - c\boldsymbol{m})_i \le e^{-c}\, e^{z_i} \big/ \sum_{j \in U} e^{z_j}$ である。

証明（命題 3.3）

$S_U := \sum_{j \in U} e^{z_j} > 0$（$U$ は空でない）、$S_{\bar U} := \sum_{j \notin U} e^{z_j} \ge 0$ と置きます。$\boldsymbol{z} - c\boldsymbol{m}$ の成分は、$i \in U$ で $z_i$、$i \notin U$ で $z_i - c$ なので、分母は $S_U + e^{-c}S_{\bar U}$ です。

$i \in U$ のとき $\sigma(\boldsymbol{z} - c\boldsymbol{m})_i = \dfrac{e^{z_i}}{S_U + e^{-c}S_{\bar U}}$。$c \to \infty$ で $e^{-c} = 1/e^{c} \to 0$（$e^{c} \ge 1 + c$。定理 3.3［指数・対数・三角関数をきちんと定義する］ の 4）なので、分母は $S_U$ に収束し、値は $e^{z_i}/S_U = \sigma^{U}(\boldsymbol{z})_i$ に収束します。

$i \notin U$ のとき $\sigma(\boldsymbol{z} - c\boldsymbol{m})_i = \dfrac{e^{-c}e^{z_i}}{S_U + e^{-c}S_{\bar U}} \le \dfrac{e^{-c}e^{z_i}}{S_U}$ で、右辺は 0 に収束します。正であることは 命題 3.1 の 1 です。

∎

実装では、マスクされた成分に十分大きな負の数（または浮動小数点の $-\infty$）を足します。上の評価から、誤差は $e^{-c}$ に比例して小さくなります。

4. スケール付き内積注意とマルチヘッド

定義 4.1（スケール付き内積注意・自己注意・マルチヘッド・因果マスク）↗

（注意） $Q \in \mathbb{R}^{n \times d_k}$、$K \in \mathbb{R}^{m \times d_k}$、$V \in \mathbb{R}^{m \times d_v}$ に対し

$$A := \sigma\!\left( \frac{Q K^{\mathsf{T}}}{\sqrt{d_k}} \right) \in \mathbb{R}^{n \times m}, \qquad \mathrm{Attn}(Q, K, V) := A\,V \in \mathbb{R}^{n \times d_v}$$

をスケール付き内積注意、$A$ を注意の重み行列という（$\sigma$ は行ごと）。$Q, K, V$ の行をそれぞれクエリ、キー、バリューという。成分で書けば、$A_{ij} = \sigma(\boldsymbol{s}_i)_j$、$s_{ij} = \langle \boldsymbol{q}_i, \boldsymbol{k}_j\rangle/\sqrt{d_k}$ で、出力の第 $i$ 行は $\sum_{j=1}^{m} A_{ij}\boldsymbol{v}_j$ である。

（自己注意） 入力 $X \in \mathbb{R}^{n \times d}$ と、パラメータ $W_Q, W_K \in \mathbb{R}^{d \times d_k}$、$W_V \in \mathbb{R}^{d \times d_v}$、$\boldsymbol{b}_Q, \boldsymbol{b}_K \in \mathbb{R}^{d_k}$、$\boldsymbol{b}_V \in \mathbb{R}^{d_v}$ に対し

$$Q = XW_Q + \boldsymbol{1}\boldsymbol{b}_Q^{\mathsf{T}}, \qquad K = XW_K + \boldsymbol{1}\boldsymbol{b}_K^{\mathsf{T}}, \qquad V = XW_V + \boldsymbol{1}\boldsymbol{b}_V^{\mathsf{T}}$$

として $\mathrm{Attn}(Q, K, V)$ を計算することを（1 ヘッドの）自己注意という。クエリ・キー・バリューがすべて同じ入力 $X$ から作られる。

（マルチヘッド） $h$ を $d$ の約数とし、$d_k = d_v = d/h$ とする。$a = 1, \ldots, h$ について別々のパラメータで 1 ヘッドの自己注意の出力 $H^{(a)} \in \mathbb{R}^{n \times d_k}$ を計算し、列方向に連結した $[H^{(1)} \cdots H^{(h)}] \in \mathbb{R}^{n \times d}$ に、パラメータ $W_O \in \mathbb{R}^{d \times d}$、$\boldsymbol{b}_O \in \mathbb{R}^{d}$ を使って

$$\mathrm{MHA}(X) := [H^{(1)} \cdots H^{(h)}]\, W_O + \boldsymbol{1}\boldsymbol{b}_O^{\mathsf{T}} \in \mathbb{R}^{n \times d}$$

と定める。これをマルチヘッド自己注意、$h$ をヘッド数という。

（因果マスク） 自己注意（$m = n$）で、重みの第 $i$ 行をマスクつき softmax $\sigma^{U_i}$、$U_i := \{1, \ldots, i\}$ で計算すること、すなわち $j > i$ の成分に $-\infty$ を足すことを因果マスクという。このとき $A_{ij} = 0$（$j > i$）である。

Vaswani らの式では射影は行列を掛けるだけですが（[1] 3.2.2 節）、ここではバイアスも含めて定義しました。バイアスを付けない場合は以下の命題で $\boldsymbol{b}$ を 0 と置けばよく、パラメータ数だけが変わります。ヘッドを複数にするのは、「主語を探すヘッド」「直前の語を見るヘッド」のように、違う基準の重みを同時に持てるようにするためです。$d_k = d/h$ と取るので、ヘッドを増やしても全体の大きさは変わりません（命題 7.1、命題 7.2）。

命題 4.2（注意の出力は凸結合）✓

定義 4.1 の注意（因果マスクつきの場合を含む）について、出力の第 $i$ 行は

$$\sum_{j=1}^{m} A_{ij}\boldsymbol{v}_j, \qquad A_{ij} \ge 0, \qquad \sum_{j=1}^{m} A_{ij} = 1$$

の形、すなわちバリュー $\boldsymbol{v}_1, \ldots, \boldsymbol{v}_m$ の凸結合である。したがって

$$\Big\| \sum_{j} A_{ij}\boldsymbol{v}_j \Big\| \le \max_{1 \le j \le m} \|\boldsymbol{v}_j\|$$

が成り立ち、各成分についても $\min_j (\boldsymbol{v}_j)_l \le \big(\sum_j A_{ij}\boldsymbol{v}_j\big)_l \le \max_j (\boldsymbol{v}_j)_l$ である。

証明（命題 4.2）

マスクなしのとき、$A$ の第 $i$ 行は $\sigma(\boldsymbol{s}_i)$ で、命題 3.1 の 1 より成分は正で和は 1 です。因果マスクつきのとき、第 $i$ 行は $\sigma^{U_i}(\boldsymbol{s}_i)$ で、$j \in U_i$ の成分は正、$j \notin U_i$ の成分は 0、和は $\sum_{j \in U_i} e^{s_{ij}} / \sum_{l \in U_i}e^{s_{il}} = 1$ です。

ノルムについて、三角不等式（系 4.3［内積空間とグラム・シュミット直交化］）と $\|c\boldsymbol{v}\| = |c|\|\boldsymbol{v}\|$（定義 4.1［内積空間とグラム・シュミット直交化］）より

$$\Big\| \sum_j A_{ij}\boldsymbol{v}_j \Big\| \le \sum_j A_{ij}\|\boldsymbol{v}_j\| \le \Big(\sum_j A_{ij}\Big)\max_j\|\boldsymbol{v}_j\| = \max_j \|\boldsymbol{v}_j\|.$$

成分については、$\mu := \min_j(\boldsymbol{v}_j)_l$ として $\sum_j A_{ij}(\boldsymbol{v}_j)_l \ge \sum_j A_{ij}\mu = \mu$。上からも同様です。

∎

注意は、バリューの集まりの「内側」しか出力できません。新しい方向を作るのは、前後の線形変換と、次章の MLP の役目です。

例 4.3（n = 3, d = 2 の手計算）

$d = d_k = d_v = 2$、1 ヘッド、$W_Q = W_K = W_V = I$、バイアス 0 とし、

$$X = \begin{pmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{pmatrix}$$

とします。$Q = K = V = X$ で、$XX^{\mathsf{T}}/\sqrt2$ は

$$\frac{1}{\sqrt2}\begin{pmatrix} 1 & 0 & 1 \\ 0 & 1 & 1 \\ 1 & 1 & 2 \end{pmatrix} = \begin{pmatrix} 0.7071 & 0 & 0.7071 \\ 0 & 0.7071 & 0.7071 \\ 0.7071 & 0.7071 & 1.4142 \end{pmatrix}$$

です。$e^{0.7071} = 2.0281$、$e^{1.4142} = 4.1133$ なので、第 1 行の重みは $(2.0281,\ 1,\ 2.0281)/5.0562 = (0.4011,\ 0.1978,\ 0.4011)$、第 3 行は $(2.0281,\ 2.0281,\ 4.1133)/8.1695 = (0.2483,\ 0.2483,\ 0.5035)$ です。

$$A = \begin{pmatrix} 0.4011 & 0.1978 & 0.4011 \\ 0.1978 & 0.4011 & 0.4011 \\ 0.2483 & 0.2483 & 0.5035 \end{pmatrix}, \qquad AX = \begin{pmatrix} 0.8022 & 0.5989 \\ 0.5989 & 0.8022 \\ 0.7517 & 0.7517 \end{pmatrix}$$

出力の第 1 行は $0.4011\,(1, 0) + 0.1978\,(0, 1) + 0.4011\,(1, 1) = (0.8022,\ 0.5989)$ で、3 つの行ベクトルが作る三角形の内側の点です（命題 4.2）。

行の入れ替え。 第 1 行と第 2 行を入れ替えた入力 $\big((0,1), (1,0), (1,1)\big)$ で同じ計算をすると、出力は $\big((0.5989, 0.8022),\ (0.8022, 0.5989),\ (0.7517, 0.7517)\big)$ になります。もとの出力の第 1 行と第 2 行が入れ替わっただけで、第 3 行は変わりません。これが 5 節の同変性です。

因果マスク。 マスクを入れると、第 1 行は自分だけを見るので重みは $(1, 0, 0)$、第 2 行は $(1,\ 2.0281)/3.0281 = (0.3302,\ 0.6698)$ と 0、第 3 行は変わりません。出力は $\big((1, 0),\ (0.3302, 0.6698),\ (0.7517, 0.7517)\big)$ です。ここで第 3 行を $(5, -2)$ に変えると、出力は小数第 4 位まででは $\big((1, 0),\ (0.3302, 0.6698),\ (5.0000, -2.0000)\big)$ になり、第 1・2 行は変わりません（第 3 行は自分自身との内積 $29/\sqrt2$ が大きく、重みがほぼ $(0, 0, 1)$ になります）。マスクなしで同じ変更をすると第 1 行は $(4.6488, -1.8110)$ に変わります。これが 6 節の因果性です。

4.1. なぜ $\sqrt{d_k}$ で割るのか

Vaswani らは、内積を $\sqrt{d_k}$ で割る理由を脚注で説明しています。クエリとキーの成分が独立で平均 0・分散 1 なら、内積の分散は $d_k$ になり、次元が大きいほど softmax の入力が散らばる、というものです（[1] 3.2.1 節）。この主張を証明します。本筋は、独立な項の分散は足し合わせられるので、$d_k$ 個の項の合計の分散が $d_k$ になる、という一点です。証明の中の「$q_i^2$ と $k_i^2$ も独立」といった確認は、その足し算を使ってよいことの裏づけなので、初めは結論だけ追っても構いません。

命題 4.4（内積の分散は次元に比例する）✓

$q_1, \ldots, q_{d_k}, k_1, \ldots, k_{d_k}$ を独立な確率変数とし、どれも平均 0・分散 1 とする。$\boldsymbol{q} = (q_1, \ldots, q_{d_k})$、$\boldsymbol{k} = (k_1, \ldots, k_{d_k})$ について

$$E[\langle \boldsymbol{q}, \boldsymbol{k}\rangle] = 0, \qquad \mathrm{Var}\big(\langle \boldsymbol{q}, \boldsymbol{k}\rangle\big) = d_k, \qquad \mathrm{Var}\!\left(\frac{\langle \boldsymbol{q}, \boldsymbol{k}\rangle}{\sqrt{d_k}}\right) = 1$$

が成り立つ。

証明（命題 4.4）

独立な族の一部を取り出したものも独立です（定義 2.3［確率変数と期待値］ で、取り出さない変数に対応するボレル集合を $\mathbb{R}$ に取る）。分散が有限なので各変数は 2 乗可積分で、したがって可積分です。

平均。 $q_i$ と $k_i$ は独立で可積分なので、命題 7.5［確率変数と期待値］ より $E[q_ik_i] = E[q_i]E[k_i] = 0$。期待値の線形性より $E[\langle\boldsymbol{q},\boldsymbol{k}\rangle] = \sum_i E[q_ik_i] = 0$。

各項の分散。 $q_i^2$ と $k_i^2$ も独立です。実際、ボレル集合 $B$ に対し $\{q_i^2 \in B\} = \{q_i \in g^{-1}(B)\}$（$g(t) = t^2$）で、$g$ は連続なので $g^{-1}(B)$ はボレル集合であり、独立性の定義の式がそのまま使えます。$E[q_i^2] = \mathrm{Var}(q_i) + (E[q_i])^2 = 1$（命題 5.3［確率変数と期待値］ の 1）で、$k_i$ も同様なので、再び 命題 7.5［確率変数と期待値］ より $E[(q_ik_i)^2] = E[q_i^2]E[k_i^2] = 1$。よって $q_ik_i$ は 2 乗可積分で、$\mathrm{Var}(q_ik_i) = E[(q_ik_i)^2] - (E[q_ik_i])^2 = 1$。

共分散。 $i \ne j$ のとき、$q_i, k_i, q_j, k_j$ は独立な 4 個の可積分な確率変数なので、補題 10.5（Appendix）より $E[q_ik_iq_jk_j] = E[q_i]E[k_i]E[q_j]E[k_j] = 0$。よって $\mathrm{Cov}(q_ik_i, q_jk_j) = E[q_ik_iq_jk_j] - E[q_ik_i]E[q_jk_j] = 0$。

合計。 命題 5.3［確率変数と期待値］ の 4 より

$$\mathrm{Var}\Big(\sum_{i=1}^{d_k} q_ik_i\Big) = \sum_{i=1}^{d_k}\mathrm{Var}(q_ik_i) + 2\sum_{i<j}\mathrm{Cov}(q_ik_i, q_jk_j) = d_k .$$

同じ命題の 2 より $\mathrm{Var}\big(\langle\boldsymbol{q},\boldsymbol{k}\rangle/\sqrt{d_k}\big) = \mathrm{Var}(\langle\boldsymbol{q},\boldsymbol{k}\rangle)/d_k = 1$ です。

∎

例 4.5（割らないと重みが 1 点に集中する）

標準正規分布から成分を引いて確かめます（コードは 8 節）。左の 2 列は 20,000 組の $(\boldsymbol{q}, \boldsymbol{k})$ での標本分散、右の 2 列は、1 個のクエリと 8 個のキーで作った重みの最大値の平均（2,000 回）です。

$d_k$	$\langle\boldsymbol{q},\boldsymbol{k}\rangle$ の分散	$\sqrt{d_k}$ で割ったあとの分散	最大の重み（割らない）	最大の重み（割る）
4	4.00	1.001	0.532	0.347
16	16.18	1.011	0.743	0.353
64	62.92	0.983	0.878	0.361
256	255.84	0.999	0.939	0.366
1024	1036.38	1.012	0.971	0.361

分散は 命題 4.4 のとおり $d_k$ に比例します。割らないと、次元が大きいほど 8 個のうち 1 個に重みが集中し、$d_k = 1024$ では最大の重みが平均 0.97 です。命題 3.2 の 3 より、そのときヤコビ行列の成分は $1 - 0.97 = 0.03$ 程度以下で、勾配が流れにくくなります。$\sqrt{d_k}$ で割ると、最大の重みは次元に依らず 0.35 前後に保たれます。

なお、ここでの「成分が独立で分散 1」は初期化の直後を想定した模型で、学習が進んだあとのクエリとキーがこの仮定を満たすわけではありません。単位球面上の一様な方向どうしの内積の分散が $1/d$ になること（命題 3.2［数学を学び直す意義］）は、長さを 1 にそろえた場合の同じ現象です。

5. 置換同変性：自己注意は語順を知らない

この節では、「入力の行を並べ替えてから自己注意に通しても、通してから出力の行を並べ替えても同じ」ことを示します。2 節の例でいえば、1 番目と 2 番目の語を入れ替えて計算しても、出力の 1 番目と 2 番目が入れ替わるだけです。つまり自己注意は語がどの順に並んでいるかを知りません。まず、この性質を保ったまま部品を組み立てられることを補題で確かめ、そのあとで自己注意そのものを調べます。

補題 5.1（置換同変な写像の組み立て）✓

1. 行ごとに働く写像は置換同変である。とくに $X \mapsto XW + \boldsymbol{1}\boldsymbol{b}^{\mathsf{T}}$ は置換同変である。
2. $f, g$ が置換同変ならば、合成 $g \circ f$、和 $X \mapsto f(X) + g(X)$、列方向の連結 $X \mapsto [f(X)\ g(X)]$ も置換同変である。とくに $f$ が置換同変ならば $X \mapsto X + f(X)$ も置換同変である。

証明（補題 5.1）

$P = P_\pi$ とします。$PX$ の第 $i$ 行は $\boldsymbol{x}_{\pi(i)}$ です。

1. $G(PX)_i = g\big((PX)_i\big) = g(\boldsymbol{x}_{\pi(i)}) = G(X)_{\pi(i)} = (PG(X))_i$。$XW + \boldsymbol{1}\boldsymbol{b}^{\mathsf{T}}$ の第 $i$ 行は $W^{\mathsf{T}}\boldsymbol{x}_i + \boldsymbol{b}$ で、$g(\boldsymbol{x}) = W^{\mathsf{T}}\boldsymbol{x} + \boldsymbol{b}$ を各行に施したものです。
2. $g(f(PX)) = g(Pf(X)) = Pg(f(X))$。$f(PX) + g(PX) = Pf(X) + Pg(X) = P(f(X) + g(X))$。連結は、$P$ を左から掛けることが行の並べ替えで、連結した行列の第 $i$ 行は $f(X)$ と $g(X)$ の第 $i$ 行を並べたものなので、$[f(PX)\ g(PX)] = [Pf(X)\ Pg(X)] = P[f(X)\ g(X)]$。最後の主張は、恒等写像が置換同変であることと和の場合から従います。

∎

定理 5.2（自己注意の置換同変性）✓

位置情報を加えておらず、マスクもないマルチヘッド自己注意 $\mathrm{MHA} : \mathbb{R}^{n\times d} \to \mathbb{R}^{n \times d}$ は置換同変である。すなわち任意の置換行列 $P$ と任意の $X$ について

$$\mathrm{MHA}(PX) = P\,\mathrm{MHA}(X).$$

さらに、マスクのない自己注意、行ごとに働く写像（MLP・正規化など）、それらの合成・和（残差接続）だけから組み立てた写像は、すべて置換同変である。

証明（定理 5.2）

$P = P_\pi$ とし、1 つのヘッドについて、入力を $PX$ に替えたときの量にダッシュを付けます。

クエリ・キー・バリュー。 積の結合法則（系 4.3［行列と連立一次方程式］）と $P\boldsymbol{1} = \boldsymbol{1}$ より

$$Q' = (PX)W_Q + \boldsymbol{1}\boldsymbol{b}_Q^{\mathsf{T}} = P(XW_Q) + P\boldsymbol{1}\boldsymbol{b}_Q^{\mathsf{T}} = PQ$$

で、同様に $K' = PK$、$V' = PV$ です。

点数。 $S := QK^{\mathsf{T}}/\sqrt{d_k}$ と置くと、$S' = Q'K'^{\mathsf{T}}/\sqrt{d_k} = PQK^{\mathsf{T}}P^{\mathsf{T}}/\sqrt{d_k} = PSP^{\mathsf{T}}$。2 節の事実より $S'_{ij} = S_{\pi(i)\pi(j)}$ です。

重み。 $S'$ の第 $i$ 行は、$S$ の第 $\pi(i)$ 行 $\boldsymbol{s}_{\pi(i)}$ の成分を $\pi$ で並べ替えたベクトル $P\boldsymbol{s}_{\pi(i)}$ です。softmax の置換同変性（命題 3.1 の 4）より $\sigma(P\boldsymbol{s}_{\pi(i)}) = P\sigma(\boldsymbol{s}_{\pi(i)})$ で、その第 $j$ 成分は $A_{\pi(i)\pi(j)}$。よって $A'_{ij} = A_{\pi(i)\pi(j)}$、すなわち $A' = PAP^{\mathsf{T}}$ です。

出力。 $P^{\mathsf{T}}P = I$ より

$$A'V' = PAP^{\mathsf{T}}PV = PAV .$$

よって各ヘッドの出力は $H'^{(a)} = PH^{(a)}$ で、1 ヘッドの自己注意は置換同変です。

マルチヘッド。 $\mathrm{MHA}$ は、置換同変な $H^{(1)}, \ldots, H^{(h)}$ を連結し、$Y \mapsto YW_O + \boldsymbol{1}\boldsymbol{b}_O^{\mathsf{T}}$ を合成したものなので、補題 5.1 より置換同変です。後半の主張も同じ補題から、組み立ての段数についての帰納法で従います。

∎

証明の要点は、点数の行列が $PSP^{\mathsf{T}}$ と「行と列を同時に」並べ替わり、バリューの並べ替え $P$ と $P^{\mathsf{T}}P = I$ で打ち消し合うことです。この定式化は、集合を入力とするネットワークの文脈で Lee らが整理しています（[3]）。

系 5.3（位置情報がなければ語順を区別できない）✓

$f : \mathbb{R}^{n\times d} \to \mathbb{R}^{n\times d'}$ を置換同変な写像とする（定理 5.2 より、位置情報もマスクもない Transformer の層の重ね合わせはこれにあたる）。

1. 置換 $\pi$ が $\pi(i) = i$ を満たすならば $f(P_\pi X)_i = f(X)_i$。すなわち出力の第 $i$ 行は、$\boldsymbol{x}_i$ と、ほかの行の多重集合（順序を忘れ、重複は数える集まり）だけで決まる。
2. とくに最後の位置 $i = n$ の出力は、それより前のトークンをどう並べ替えても変わらない。

証明（系 5.3）

1. 同変性より $f(P_\pi X)_i = (P_\pi f(X))_i = f(X)_{\pi(i)} = f(X)_i$。$i$ 以外の行の並べ方が違う 2 つの入力は、$i$ を動かさない置換で移り合うので、第 $i$ 行の出力は等しくなります。2 は 1 で $i = n$ とした場合です。

∎

次のトークンの予測は最後の位置の出力から作るので、2 は「位置情報のないモデルは、『犬が猫を追う』と『猫が犬を追う』の続きについて同じ予測をする」ことを意味します（トークン化して最後のトークンが同じなら）。前章の位置符号化を入れると、$X$ の各行に位置ごとに違うベクトルが足されるので、行を入れ替えた入力はもはや「同じ行の並べ替え」ではなくなり、この系の仮定が外れます。RoPE（回転位置埋め込み。前章で扱いました）の場合は点数 $S_{ij}$ が $j - i$ に依るので $S' = PSP^{\mathsf{T}}$ が成り立たなくなります。

注意 5.4（因果マスクつきの場合）

因果マスクは「$j \le i$ だけを見る」という形で位置の添字を使うので、因果マスクつきの自己注意は置換同変ではありません（8 節の数値実験で、$f(PX)$ と $Pf(X)$ が大きく違うことを確かめます）。ただし、マスクが位置の情報を十分に伝えるかどうかは別の問題です。演習 10.4 では、因果マスクつきの注意 1 層の最後の行が、前の行の並べ替えで変わらないことを示します。層を重ねた場合にどこまで語順を使えるかは、この章では扱いません。

6. 因果マスク：未来に依存しない

言語モデルは「これまでのトークンから次のトークンを予測する」関数です。位置 $i$ の出力が $i$ より後ろのトークンに依存してしまうと、学習のときに答えを見て予測することになります。因果マスクはそれを防ぎます。入力の長さ $n$ を変えたときの振る舞いまで含めて述べるために、次の定義を置きます。

次の定義の式 $f_n(X)_i = f_i(X_{1:i})_i$ は、「長さ $n$ の文全体を入れて第 $i$ 行を見ても、最初の $i$ 語だけを入れて最後の行を見ても、答えは同じ」と読んでください。たとえば $n = 5$、$i = 3$ なら、5 語を入れたときの 3 行目は、3 語だけを入れたときの 3 行目（最後の行）と等しい、ということです。

定義 6.1（因果的な写像の族）↗

各 $n \ge 1$ に対して写像 $f_n : \mathbb{R}^{n \times d} \to \mathbb{R}^{n \times d'}$ が与えられているとする。族 $f = (f_n)_{n \ge 1}$ が因果的であるとは、すべての $n$、すべての $i \le n$、すべての $X \in \mathbb{R}^{n \times d}$ について

$$f_n(X)_i = f_i(X_{1:i})_i$$

が成り立つことをいう。ここで $X_{1:i} \in \mathbb{R}^{i \times d}$ は $X$ の最初の $i$ 行からなる行列である。

因果的ならば、$f_n(X)$ の第 $i$ 行は $\boldsymbol{x}_1, \ldots, \boldsymbol{x}_i$ だけの関数です（右辺が $X_{1:i}$ にしか依らないため）。それに加えて、「長さ $n$ の列の第 $i$ 行の出力」と「最初の $i$ 個だけを入力したときの最後の行の出力」が一致することも言っています。自己注意は同じパラメータで任意の長さの列に適用できるので、族 $(f_n)$ が自然に定まります。

補題 6.2（因果的な写像の組み立て）✓

1. 関数 $g_1, g_2, \ldots : \mathbb{R}^d \to \mathbb{R}^{d'}$（$n$ に依らない）に対し、$G_n(X)_i := g_i(\boldsymbol{x}_i)$ で定まる族は因果的である。とくに、行ごとに働く写像（すべての $g_i$ が同じ場合）と、位置 $i$ に依る位置ベクトルを足す写像 $\boldsymbol{x}_i \mapsto \boldsymbol{x}_i + \boldsymbol{p}_i$ は因果的である。
2. $f, g$ が因果的ならば、合成 $(g_n \circ f_n)_n$、和 $(f_n + g_n)_n$、列方向の連結も因果的である。とくに $f$ が因果的ならば $X \mapsto X + f_n(X)$ も因果的である。
3. $f$ が因果的ならば、$j \le i \le n$ について $f_n(X)_j = f_i(X_{1:i})_j$。すなわち $f_n(X)_{1:i} = f_i(X_{1:i})$。

証明（補題 6.2）

1. $G_n(X)_i = g_i(\boldsymbol{x}_i)$ で、$X_{1:i}$ の第 $i$ 行も $\boldsymbol{x}_i$ なので $G_i(X_{1:i})_i = g_i(\boldsymbol{x}_i)$。両者は等しい。

2. $j \le i$ とします。定義を $(n, j)$ に使うと $f_n(X)_j = f_j(X_{1:j})_j$。定義を $(i, j)$ と入力 $X_{1:i}$ に使うと $f_i(X_{1:i})_j = f_j\big((X_{1:i})_{1:j}\big)_j = f_j(X_{1:j})_j$。よって両者は等しい。

3. 合成について。3 より $f_n(X)_{1:i} = f_i(X_{1:i})$ なので、$g$ の因果性と合わせて

$$g_n(f_n(X))_i = g_i\big(f_n(X)_{1:i}\big)_i = g_i\big(f_i(X_{1:i})\big)_i .$$

右辺は合成の族を $X_{1:i}$ に施したものの第 $i$ 行です。和と連結は、第 $i$ 行がそれぞれの第 $i$ 行の和・連結なので、$f$ と $g$ の因果性からただちに従います。最後の主張は、恒等写像が 1 の形（$g_i$ が恒等写像）で因果的であることと和の場合からです。

∎

定理 6.3（因果マスクつきの自己注意は未来に依存しない）✓

因果マスクつきのマルチヘッド自己注意は因果的である。したがって出力の第 $i$ 行は $\boldsymbol{x}_1, \ldots, \boldsymbol{x}_i$ だけの関数である。さらに、因果マスクつきの自己注意、補題 6.2 の 1 の形の写像（行ごとの MLP・正規化、位置ベクトルの加算）、それらの合成・和（残差接続）を何層重ねた写像も因果的である。

証明（定理 6.3）

1 ヘッド。 長さ $n$ の入力 $X$ に対する出力の第 $i$ 行は、定義 4.1 より

$$\sum_{j=1}^{i} A_{ij}\boldsymbol{v}_j, \qquad A_{ij} = \frac{e^{s_{ij}}}{\sum_{l=1}^{i} e^{s_{il}}}, \qquad s_{ij} = \frac{\langle W_Q^{\mathsf{T}}\boldsymbol{x}_i + \boldsymbol{b}_Q,\ W_K^{\mathsf{T}}\boldsymbol{x}_j + \boldsymbol{b}_K\rangle}{\sqrt{d_k}}, \qquad \boldsymbol{v}_j = W_V^{\mathsf{T}}\boldsymbol{x}_j + \boldsymbol{b}_V$$

です（$j > i$ の重みは 0 なので和は $j \le i$ だけ）。この式に現れる添字は $1, \ldots, i$ だけで、$n$ も $\boldsymbol{x}_{i+1}, \ldots, \boldsymbol{x}_n$ も現れません。入力を $X_{1:i}$ に替えて第 $i$ 行（最後の行）を計算しても、まったく同じ式になります。よって $f_n(X)_i = f_i(X_{1:i})_i$ です。

マルチヘッド。 各ヘッドが因果的なので、連結と、行ごとに働く写像 $Y \mapsto YW_O + \boldsymbol{1}\boldsymbol{b}_O^{\mathsf{T}}$ との合成も因果的です（補題 6.2 の 1 と 2）。

層の重ね合わせ。 層の数 $L$ についての帰納法です。$L = 1$ は上で示した場合か、補題 6.2 の 1 の形です。$L - 1$ 層までの合成が因果的なら、それに因果的な層を合成・加算したものも 補題 6.2 の 2 より因果的です。

∎

系 6.4（1 回の順伝播で n 個の予測が得られる）✓

$f = (f_n)$ を因果的な族とする（定理 6.3 の形のモデル）。長さ $n$ の列 $X$ に対する 1 回の計算 $f_n(X)$ の第 $t$ 行は、接頭辞 $X_{1:t}$ だけを入力したときの最後の行に等しい：

$$f_n(X)_t = f_t(X_{1:t})_t \qquad (t = 1, \ldots, n).$$

したがって、$n$ 個の「接頭辞 $(\boldsymbol{x}_1, \ldots, \boldsymbol{x}_t)$ から次のトークンを予測する」計算の結果が、1 回の順伝播で同時に、別々に計算したときと同じ値で得られる。

証明（系 6.4）

定義 6.1 の定義そのものです。モデルが因果的であることは 定理 6.3 です。

∎

この系が、言語モデルの学習を効率よくしている理由です。長さ $n$ の文 1 つから、$n$ 個の予測問題の損失とその勾配が 1 回の順伝播・逆伝播で得られます。接頭辞ごとに別々に計算すれば、長さ $1, 2, \ldots, n$ の列を $n$ 回処理することになります。生成のときには逆に、補題 6.2 の 3（$f_n(X)_{1:i} = f_i(X_{1:i})$）が効きます。トークンを 1 個足しても、それまでの位置の計算結果は変わらないので、使い回せます。第 6 章の KV キャッシュはこの事実に基づきます。

n = 4 の因果マスク。行が「見る側」（クエリ）の位置 i、列が「見られる側」（キー）の位置 j。塗った成分（j ≤ i）だけ重みが正で、ほかは 0。重みの行列は下三角になる。

7. パラメータ数と計算量

命題 7.1（マルチヘッド自己注意のパラメータ数）✓

定義 4.1 のマルチヘッド自己注意（$d_k = d_v = d/h$、バイアスあり）1 層のパラメータ数は

$$4d^2 + 4d$$

で、ヘッド数 $h$ に依らない。バイアスを付けない場合は $4d^2$ である。

証明（命題 7.1）

1 つのヘッドは $W_Q, W_K, W_V \in \mathbb{R}^{d \times d_k}$ と $\boldsymbol{b}_Q, \boldsymbol{b}_K, \boldsymbol{b}_V \in \mathbb{R}^{d_k}$ を持つので $3(d\,d_k + d_k)$ 個。$h$ ヘッドで $3h(d\,d_k + d_k) = 3(d\cdot hd_k + hd_k) = 3(d^2 + d)$ 個（$hd_k = d$）。これに $W_O \in \mathbb{R}^{d\times d}$ と $\boldsymbol{b}_O \in \mathbb{R}^d$ の $d^2 + d$ 個を加えて $4d^2 + 4d$ 個です。バイアスを除けば $3d^2 + d^2 = 4d^2$ です。

∎

たとえば Vaswani らの基本の構成 $d = 512$、$h = 8$、$d_k = 64$（[1]）なら、$4 \times 512^2 + 4 \times 512 = 1{,}050{,}624$ 個（バイアスなしなら $1{,}048{,}576$ 個）です。

命題 7.2（自己注意 1 層の計算量）✓

長さ $n$、次元 $d$、ヘッド数 $h$（$d_k = d/h$）のマルチヘッド自己注意 1 層を定義の式どおりに計算するとき、スカラーの乗算の回数は

$$\underbrace{4nd^2}_{\text{射影}} + \underbrace{2n^2 d}_{\text{点数と重みつき和}}$$

で、ヘッド数 $h$ に依らない。このほかに指数関数の評価と除算が $O(hn^2)$ 回ある。注意の重みを保持するのに必要な記憶は $hn^2$ 個の数である。$h \le d$ なので、時間計算量は $O(nd^2 + n^2 d)$ であり、$n \ge d$ のときは $O(n^2 d)$ である。$n$ を 2 倍にすると、$2n^2d$ の項は 4 倍、$4nd^2$ の項は 2 倍になる。

証明（命題 7.2）

$p \times q$ 行列と $q \times r$ 行列の積の乗算の回数は $pqr$ です（定義 5.1［NumPy で書く配列計算］）。

• 射影 $XW_Q$、$XW_K$、$XW_V$：1 ヘッドあたり $3 \cdot n \cdot d \cdot d_k$、$h$ ヘッドで $3nd \cdot hd_k = 3nd^2$。
• 点数 $QK^{\mathsf{T}}$：1 ヘッドあたり $n \cdot d_k \cdot n$、$h$ ヘッドで $n^2 \cdot hd_k = n^2 d$。
• 重みつき和 $AV$：1 ヘッドあたり $n \cdot n \cdot d_k$、$h$ ヘッドで $n^2 d$。
• 出力の射影（連結したものに $W_O$ を掛ける）：$n \cdot d \cdot d = nd^2$。

合計は $4nd^2 + 2n^2d$ です。$\sqrt{d_k}$ での除算、指数関数、正規化の除算は、$A$ の成分 1 個につき定数回で、$A$ は $h$ 個の $n \times n$ 行列なので $O(hn^2)$ 回です。$h \le d$ より $hn^2 \le n^2 d$ で、全体は $O(nd^2 + n^2d)$（定義 3.1［計算量と O 記法］）。$n \ge d$ なら $nd^2 \le n^2 d$ なので $O(n^2d)$ です。記憶は $A$ の成分数 $hn^2$ です。

2 倍にしたときの比は、$2n^2d$ が $n$ の 2 次、$4nd^2$ が $n$ の 1 次の単項式であることから、それぞれ $2^2 = 4$、$2^1 = 2$ です（命題 5.1［計算量と O 記法］ の 3）。

∎

Vaswani らの表 1 は自己注意の 1 層あたりの計算量を $O(n^2 \cdot d)$ としています（[1]）。上の命題の $2n^2d$ の項がそれにあたります。$n$ が $d$ より小さいあいだは射影の $4nd^2$ のほうが大きく、列が長くなると $n$ の 2 乗の項が支配的になります。$d = 64$ のときの回数は次のとおりです（8 節のコードの出力）。

$n$	$4nd^2$	$2n^2d$	合計	$2n^2d$ の前の行との比	合計の前の行との比
64	1,048,576	524,288	1,572,864
128	2,097,152	2,097,152	4,194,304	4.00	2.67
256	4,194,304	8,388,608	12,582,912	4.00	3.00
512	8,388,608	33,554,432	41,943,040	4.00	3.33
1024	16,777,216	134,217,728	150,994,944	4.00	3.60

合計の比は、$n$ が大きくなるにつれて 4 に近づきます。$n = 2d = 128$ で 2 つの項が等しくなります。

注意 7.3（交差注意）

クエリを一方の列から、キーとバリューを別の列から作る注意を交差注意といいます。Vaswani らのエンコーダ・デコーダ型のモデルでは、デコーダの各層が、エンコーダの出力をキーとバリューにした交差注意を持ちます（[1] 3.2.3 節）。命題 4.2 はそのまま成り立ちます。このテーマで扱うデコーダのみのモデルには交差注意はありません。

8. 最小実装と数値検証

マルチヘッド自己注意を NumPy で実装します。ヘッドの軸を先頭に持つ 3 次元配列でパラメータを持ち、積は einsum（定義 5.1［NumPy で書く配列計算］）で書きます。バイアスの加算と $\sqrt{d_k}$ での除算はブロードキャスト（定義 3.1［NumPy で書く配列計算］）です。

einsum は、添字の名前で「どの軸をかけて、どの軸で足すか」を指定する記法です。たとえば "nd,hdk->hnk" は、$X$ の軸を n（位置）と d（次元）、W の軸を h（ヘッド）と d と k として、共通の d について積の和を取り、軸が h, n, k の配列を作る、という意味です。矢印の右に残らない添字は足し合わされます。

# Python 3.12.14, NumPy 2.5.2
import numpy as np

def softmax(z):
    """最後の軸に沿った softmax。最大値を引いてから指数を取る（平行移動不変性）。"""
    e = np.exp(z - z.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

def init_params(d, h, rng):
    dk = d // h
    p = {name: rng.standard_normal((h, d, dk)) / np.sqrt(d) for name in ("WQ", "WK", "WV")}
    p.update({name: 0.1 * rng.standard_normal((h, dk)) for name in ("bQ", "bK", "bV")})
    p["WO"] = rng.standard_normal((d, d)) / np.sqrt(d)
    p["bO"] = 0.1 * rng.standard_normal(d)
    return p

def self_attention(X, p, causal=False):
    """マルチヘッド自己注意。X は (n, d)。出力 (n, d) と重み (h, n, n) を返す。"""
    n, d = X.shape
    h, _, dk = p["WQ"].shape
    Q = np.einsum("nd,hdk->hnk", X, p["WQ"]) + p["bQ"][:, None, :]
    K = np.einsum("nd,hdk->hnk", X, p["WK"]) + p["bK"][:, None, :]
    V = np.einsum("nd,hdk->hnk", X, p["WV"]) + p["bV"][:, None, :]
    S = np.einsum("hik,hjk->hij", Q, K) / np.sqrt(dk)
    if causal:
        S = np.where(np.tril(np.ones((n, n), dtype=bool)), S, -np.inf)
    A = softmax(S)
    H = np.einsum("hij,hjk->hik", A, V)            # (h, n, dk)
    Hcat = H.transpose(1, 0, 2).reshape(n, h * dk)  # ヘッドの出力を連結
    return Hcat @ p["WO"] + p["bO"], A

if __name__ == "__main__":
    rng = np.random.default_rng(0)
    n, d, h = 6, 8, 2
    p = init_params(d, h, rng)
    X = rng.standard_normal((n, d))
    print("params:", sum(a.size for a in p.values()), "= 4d^2+4d =", 4 * d * d + 4 * d)

    perm = rng.permutation(n)
    Y, A = self_attention(X, p)
    Yp, _ = self_attention(X[perm], p)
    print("equivariance   max|f(PX) - P f(X)| = %.1e" % np.abs(Yp - Y[perm]).max())
    Yc, Ac = self_attention(X, p, causal=True)
    Ycp, _ = self_attention(X[perm], p, causal=True)
    print("causal, permuted max|f(PX) - P f(X)| = %.1e" % np.abs(Ycp - Yc[perm]).max())

    X2 = X.copy(); X2[4:] = rng.standard_normal((2, d))   # 5・6 番目のトークンを変える
    Yc2, _ = self_attention(X2, p, causal=True)
    Y2, _ = self_attention(X2, p)
    print("causal    rows 1-4 change: %.1e" % np.abs(Yc2[:4] - Yc[:4]).max())
    print("no mask   rows 1-4 change: %.1e" % np.abs(Y2[:4] - Y[:4]).max())
    pre = max(np.abs(self_attention(X[:t], p, causal=True)[0][t - 1] - Yc[t - 1]).max() for t in range(1, n + 1))
    print("prefix consistency: %.1e" % pre)

    Vh = np.einsum("nd,hdk->hnk", X, p["WV"]) + p["bV"][:, None, :]
    Hh = np.einsum("hij,hjk->hik", Ac, Vh)
    print("weights >= 0:", bool((Ac >= 0).all()), " row sums - 1: %.1e" % np.abs(Ac.sum(-1) - 1).max(),
          " upper triangle max: %.1e" % np.abs(np.triu(Ac, 1)).max())
    print("max row norm of head output %.4f <= max row norm of V %.4f"
          % (np.linalg.norm(Hh, axis=-1).max(), np.linalg.norm(Vh, axis=-1).max()))

params: 288 = 4d^2+4d = 288
equivariance   max|f(PX) - P f(X)| = 2.2e-16
causal, permuted max|f(PX) - P f(X)| = 2.6e+00
causal    rows 1-4 change: 0.0e+00
no mask   rows 1-4 change: 1.5e+00
prefix consistency: 1.1e-16
weights >= 0: True  row sums - 1: 1.1e-16  upper triangle max: 0.0e+00
max row norm of head output 2.3391 <= max row norm of V 3.5719

出力を上から読みます。

• 1 行目：パラメータの総数が 命題 7.1 の $4d^2 + 4d = 288$ に一致。
• 2 行目：マスクなしでは $f(PX) = Pf(X)$ が丸め誤差の範囲で成り立つ（定理 5.2）。
• 3 行目：因果マスクつきでは成り立たない（注意 5.4）。
• 4・5 行目：5・6 番目の入力を変えても、因果マスクつきなら 1〜4 番目の出力はまったく変わらず、マスクなしでは変わる（定理 6.3）。
• 6 行目：接頭辞だけを入力したときの最後の行が、全体を入力したときの対応する行に一致（系 6.4）。
• 7・8 行目：重みが非負で行の和が 1、上三角が 0、ヘッドの出力の行のノルムがバリューの行のノルムの最大値以下（命題 4.2）。

例 4.5 の表と 7 節の回数の表は、次のコードの出力です。

# Python 3.12.14, NumPy 2.5.2。attention.py は前のコード
import numpy as np
from attention import softmax

rng = np.random.default_rng(0)
print("  d_k  Var(q.k)  Var(q.k/sqrt(d_k))  max weight: raw  scaled")
for dk in (4, 16, 64, 256, 1024):
    q, k = rng.standard_normal((20000, dk)), rng.standard_normal((20000, dk))
    s = (q * k).sum(axis=1)
    q1, k8 = rng.standard_normal((2000, 1, dk)), rng.standard_normal((2000, 8, dk))
    scores = (q1 * k8).sum(axis=-1)                       # 1 個のクエリと 8 個のキー
    raw = softmax(scores).max(axis=-1).mean()
    scaled = softmax(scores / np.sqrt(dk)).max(axis=-1).mean()
    print("%5d %9.2f %12.3f %18.3f %7.3f" % (dk, s.var(), (s / np.sqrt(dk)).var(), raw, scaled))

d = 64
print("    n   4nd^2      2n^2d      total   ratio(2n^2d)  ratio(total)")
prev = None
for n in (64, 128, 256, 512, 1024):
    proj, attn = 4 * n * d * d, 2 * n * n * d
    if prev is None:
        print("%5d %8d %10d %10d" % (n, proj, attn, proj + attn))
    else:
        print("%5d %8d %10d %10d %10.2f %12.2f" % (n, proj, attn, proj + attn, attn / prev[0], (proj + attn) / prev[1]))
    prev = (attn, proj + attn)

  d_k  Var(q.k)  Var(q.k/sqrt(d_k))  max weight: raw  scaled
    4      4.00        1.001              0.532   0.347
   16     16.18        1.011              0.743   0.353
   64     62.92        0.983              0.878   0.361
  256    255.84        0.999              0.939   0.366
 1024   1036.38        1.012              0.971   0.361
    n   4nd^2      2n^2d      total   ratio(2n^2d)  ratio(total)
   64  1048576     524288    1572864
  128  2097152    2097152    4194304       4.00         2.67
  256  4194304    8388608   12582912       4.00         3.00
  512  8388608   33554432   41943040       4.00         3.33
 1024 16777216  134217728  150994944       4.00         3.60

9. 次章へ

この章で確定したことは 3 つです。自己注意の出力はバリューの凸結合であること、位置情報がなければ置換同変で語順を区別できないこと、因果マスクを入れれば各位置の出力が未来に依存しないこと。次章の Transformer ブロックでは、自己注意に残差接続（入力に層の出力を足す接続）・層正規化・行ごとの MLP を組み合わせて 1 つの層にし、それを重ねてモデル全体を作ります。補題 5.1 と 補題 6.2 は「行ごとに働く写像との合成と和で性質が保たれる」ことを言っているので、同変性と因果性はブロック全体にそのまま引き継がれます。

10. 演習

演習 10.1易

$\boldsymbol{z} = (0, \log 2, \log 3)$ について $\sigma(\boldsymbol{z})$ を求めよ。また $\boldsymbol{z}' = (100, 100 + \log 2, 100 + \log 3)$ について $\sigma(\boldsymbol{z}')$ を求めよ。

解答

$e^{0} = 1$、$e^{\log 2} = 2$、$e^{\log 3} = 3$ で和は 6 なので $\sigma(\boldsymbol{z}) = (1/6, 1/3, 1/2)$ です。$\boldsymbol{z}' = \boldsymbol{z} + 100\cdot\boldsymbol{1}$ なので、平行移動不変性（命題 3.1 の 2）より $\sigma(\boldsymbol{z}') = (1/6, 1/3, 1/2)$ です。定義どおりに $e^{100}$ を計算する必要はありません。

演習 10.2標準

$n = 2$ のとき、$\sigma(\boldsymbol{z}) = (s, 1 - s)$ と書く。ヤコビ行列を $s$ で表し、固有値を求めよ。固有値の最大値が最も大きくなるのは $s$ がいくつのときか。

解答

命題 3.2 より

$$J = \begin{pmatrix} s(1-s) & -s(1-s) \\ -s(1-s) & (1-s)s \end{pmatrix} = s(1-s)\begin{pmatrix} 1 & -1 \\ -1 & 1 \end{pmatrix}.$$

$\begin{pmatrix} 1 & -1 \\ -1 & 1\end{pmatrix}$ は $(1, 1)^{\mathsf{T}}$ を $\boldsymbol{0}$ に、$(1, -1)^{\mathsf{T}}$ を $2(1,-1)^{\mathsf{T}}$ に写すので、固有値は 0 と 2 です。よって $J$ の固有値は $0$（固有ベクトル $\boldsymbol{1}$。命題の 2）と $2s(1-s)$ です。$2s(1-s) = \tfrac12 - 2(s - \tfrac12)^2$ なので、$s = 1/2$ のとき最大値 $1/2$ を取り、$s$ が 0 か 1 に近づくと 0 に近づきます。出力が一方に偏るほど勾配が流れにくくなることが、この場合は式で見えます。

演習 10.3易

$\boldsymbol{z} = (1, 2, 0.5)$、$U = \{1, 2\}$ とする。$\sigma^{U}(\boldsymbol{z})$ を求めよ。また、第 3 成分から $c = 10$ を引いて普通の softmax を取ったとき、第 3 成分の値はどれだけ以下か。

解答

$e^{1} = 2.7183$、$e^{2} = 7.3891$ で和は $10.1073$ なので、$\sigma^{U}(\boldsymbol{z}) = (0.2689,\ 0.7311,\ 0)$ です。命題 3.3 の評価より、第 3 成分は $e^{-10}e^{0.5}/10.1073 = 4.54\times10^{-5} \times 1.6487 / 10.1073 \approx 7.4 \times 10^{-6}$ 以下です。

演習 10.4標準

位置情報のない、因果マスクつきの自己注意 1 層（1 ヘッド）を $f_n$ とする。$\pi$ が $\pi(n) = n$ を満たす置換のとき、$f_n(P_\pi X)_n = f_n(X)_n$ を示せ。2 層重ねたときにも同じことが成り立つか、数値実験の結果から答えよ。

解答

定理 6.3 の証明の式で $i = n$ とすると、最後の行の出力は

$$\sum_{j=1}^{n} \frac{e^{s_{nj}}}{\sum_{l=1}^{n}e^{s_{nl}}}\, \boldsymbol{v}_j$$

で、$s_{nj}$ は $\boldsymbol{x}_n$ と $\boldsymbol{x}_j$ だけ、$\boldsymbol{v}_j$ は $\boldsymbol{x}_j$ だけで決まります。入力を $P_\pi X$ に替えると、第 $j$ 行は $\boldsymbol{x}_{\pi(j)}$、最後の行は $\boldsymbol{x}_{\pi(n)} = \boldsymbol{x}_n$ のままなので、和の第 $j$ 項がもとの和の第 $\pi(j)$ 項になります。$\pi$ は全単射なので、分母の和も全体の和も、項の順番が入れ替わるだけで値は同じです。

2 層では成り立つとは限りません。1 層目の第 $j$ 行（$j < n$）の出力は $\boldsymbol{x}_1, \ldots, \boldsymbol{x}_j$ に依るので、前の行を並べ替えると、2 層目に入る行の集まりそのものが変わるからです。8 節の attention.py を使った次の実験では、1 番目と 2 番目の行を入れ替えたとき、1 層では最後の行が変わらず、2 層では変わりました。

# Python 3.12.14, NumPy 2.5.2。attention.py は前のコード
import numpy as np
from attention import init_params, self_attention

rng = np.random.default_rng(0)
n, d, h = 6, 8, 2
p1, p2 = init_params(d, h, rng), init_params(d, h, rng)
X = rng.standard_normal((n, d))
Xs = X.copy(); Xs[[0, 1]] = X[[1, 0]]          # 1 番目と 2 番目を入れ替える（最後の行は動かさない）

def one(X):
    return self_attention(X, p1, causal=True)[0]

def two(X):
    return self_attention(one(X), p2, causal=True)[0]

print("1 layer, last row change: %.1e" % np.abs(one(Xs)[-1] - one(X)[-1]).max())
print("2 layers, last row change: %.1e" % np.abs(two(Xs)[-1] - two(X)[-1]).max())

1 layer, last row change: 0.0e+00
2 layers, last row change: 5.1e-01

参考文献

• [1] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, I. Polosukhin, “Attention Is All You Need”, 2017. arXiv:1706.03762 — 3.2 節（スケール付き内積注意、マルチヘッド、因果マスク。$\sqrt{d_k}$ で割る理由は 3.2.1 節の脚注）、表 1（層の種類ごとの計算量）。
• [2] D. Bahdanau, K. Cho, Y. Bengio, “Neural Machine Translation by Jointly Learning to Align and Translate”, ICLR 2015. arXiv:1409.0473 — 注意機構を機械翻訳に導入した論文。
• [3] J. Lee, Y. Lee, J. Kim, A. R. Kosiorek, S. Choi, Y. W. Teh, “Set Transformer: A Framework for Attention-based Permutation-Invariant Neural Networks”, ICML 2019. arXiv:1810.00825 — 注意に基づく置換同変・置換不変なネットワークの定式化。
• [4] 本サイト「確率変数と期待値：可測関数としての確率変数、ルベーグ積分としての期待値」 — 独立性、分散の性質、独立な確率変数の積の期待値。

Appendix: 独立な確率変数の積の期待値

命題 4.4 の証明で使った、3 個以上の積についての補題を示します。 2 個の場合は既出の命題で、ここではそれを個数についての帰納法で延長します。

補題 10.5（独立な可積分確率変数の積の期待値（m 個））✓

$m \ge 2$ とし、$X_1, \ldots, X_m$ を独立な確率変数で、どれも可積分とする。このとき積 $X_1 X_2 \cdots X_m$ も可積分で

$$E[X_1 X_2 \cdots X_m] = E[X_1]\,E[X_2]\cdots E[X_m]$$

が成り立つ。

証明（補題 10.5）

$m$ についての帰納法です。$m = 2$ は 命題 7.5［確率変数と期待値］ です。$m \ge 3$ とし、$m - 1$ 個の場合を仮定します。

$X_2, \ldots, X_m$ の積。 $X_2, \ldots, X_m$ は独立です（独立性の定義で $B_1 = \mathbb{R}$ と取る）。帰納法の仮定より $Y := X_2\cdots X_m$ は可積分で $E[Y] = E[X_2]\cdots E[X_m]$ です。$Y$ が確率変数であることは 命題 3.4［確率変数と期待値］ です。

$X_1$ と $Y$ の独立性。 ボレル集合 $B \subset \mathbb{R}$ を固定し、$\mathbb{R}^{m-1}$ のボレル集合 $C$ に対して

$$\mu(C) := P\big(X_1 \in B,\ (X_2, \ldots, X_m) \in C\big), \qquad \nu(C) := P(X_1 \in B)\,P\big((X_2, \ldots, X_m) \in C\big)$$

と置きます。これが定義できること、すなわち任意のボレル集合 $C$ について $\{(X_2, \ldots, X_m) \in C\}$ が事象であることを先に確かめます。この性質を持つ $C$ の全体は、逆像が補集合と可算和を保つので σ-加法族で、長方形 $B_2 \times \cdots \times B_m$ を含みます（逆像は事象 $\{X_i \in B_i\}$ の共通部分）。長方形は $\mathcal{B}(\mathbb{R}^{m-1})$ を生成するので（下で使うのと同じ 命題 2.3［直積測度と Fubini–Tonelli の定理］）、すべてのボレル集合がこの性質を持ちます。$\mu, \nu$ はどちらも $\mathcal{B}(\mathbb{R}^{m-1})$ 上の有限測度です（$P$ の可算加法性から）。長方形 $B_2 \times \cdots \times B_m$（各 $B_i$ はボレル集合）の全体を $\mathcal{P}$ とすると、2 つの長方形の共通部分は長方形なので $\mathcal{P}$ は π-系で、$\mathbb{R}^{m-1} \in \mathcal{P}$ であり、$\mathcal{P}$ は $\mathcal{B}(\mathbb{R}^{m-1})$ を生成します（命題 2.3［直積測度と Fubini–Tonelli の定理］）。長方形の上では、$X_1, \ldots, X_m$ の独立性と $X_2, \ldots, X_m$ の独立性から

$$\mu(B_2 \times\cdots\times B_m) = P(X_1 \in B)\prod_{i=2}^{m}P(X_i \in B_i) = P(X_1 \in B)\,P(X_2 \in B_2, \ldots, X_m \in B_m) = \nu(B_2\times\cdots\times B_m)$$

です。測度の一意性（定理 4.5［測度空間と Carathéodory の拡張定理］）より、$\mathcal{B}(\mathbb{R}^{m-1})$ の全体で $\mu = \nu$ です。

ボレル集合 $B' \subset \mathbb{R}$ に対し $C := \{(x_2, \ldots, x_m) : x_2\cdots x_m \in B'\}$ と置くと、積 $(x_2, \ldots, x_m) \mapsto x_2\cdots x_m$ は連続なので $C$ はボレル集合で、$\{Y \in B'\} = \{(X_2, \ldots, X_m) \in C\}$ です。よって

$$P(X_1 \in B,\ Y \in B') = \mu(C) = \nu(C) = P(X_1 \in B)\,P(Y \in B')$$

で、$X_1$ と $Y$ は独立です。

結論。 $X_1$ と $Y$ は独立で可積分なので、命題 7.5［確率変数と期待値］ より $X_1Y$ は可積分で $E[X_1Y] = E[X_1]E[Y] = E[X_1]E[X_2]\cdots E[X_m]$ です。

∎