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

> Transformer の中心の演算である自己注意を定義し、softmax の基本性質とヤコビ行列、注意の出力が凸結合であること、位置情報なしでは置換同変であること、因果マスクのもとで各位置の出力が未来に依存しないこと、計算量とパラメータ数を証明し、NumPy の実装で確かめる。
> https://rikai.mugen-giken.com/computer-science/llm-architecture/self-attention

## 0. この記事の要点

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

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

<Aside type="tip" title="この記事で使う事前知識">

- **行列の積**: $p \times q$ 行列と $q \times r$ 行列から $p \times r$ 行列を作る演算で、成分は行と列の対応する成分の積の和。→ [行列の積](/mathematics/linear-algebra/matrices-and-linear-systems#def-product)
- **期待値**: 確率変数の値を確率で重みづけて平均したもの。→ [期待値](/mathematics/probability/random-variables#def-expectation)
- **分散**: 平均からのずれの 2 乗の期待値で、値の散らばりの大きさを表す。→ [分散](/mathematics/probability/random-variables#def-variance-covariance)
- **指数関数**: $e^{x}$ のこと。常に正で、$x$ が増えると増える。→ [指数関数](/mathematics/hs-analysis/elementary-functions#def-exp)
- **位置符号化**: トークンの並び順の情報をベクトルに持たせる仕組み。前章で扱いました。→ [位置符号化](/computer-science/llm-architecture/embeddings-and-positional-encoding#def-position)

</Aside>

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

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

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

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

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

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

行列 $X \in \mathbb{R}^{n \times d}$ の第 $i$ 行を列ベクトルとみなしたものを $\boldsymbol{x}_i \in \mathbb{R}^{d}$ と書きます。全成分が 1 のベクトルを $\boldsymbol{1}$ と書きます。行列の積（<Ref to="mathematics/linear-algebra/matrices-and-linear-systems#def-product" />）について、$XW$ の第 $i$ 行は $W^{\mathsf{T}}\boldsymbol{x}_i$ です。

<Definition id="def-permutation-matrix" title="置換行列・同変・不変">

$\{1, \ldots, n\}$ の置換 $\pi$（<Ref to="mathematics/linear-algebra/determinants#def-sign" />）に対し、$(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$ 行）。

</Definition>

イメージをつかむために小さい例を置きます。$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}$ です。

<Definition id="def-softmax" title="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}$ のことを指すものとする。

</Definition>

$-\infty$ は実数ではないので、「足す」という書き方はそのままでは定義になりません。上では $\sigma^U$ を直接定義し、それが極限として正当化できることを <Ref to="prop-mask-as-limit" /> で示します。

## 3. softmax の性質

<Proposition id="prop-softmax-basic" title="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})$。

</Proposition>

<Proof of="prop-softmax-basic">

$S := \sum_{j} e^{z_j}$ と置きます。指数関数は正の値を取り、狭義単調増加で、$e^{x+y} = e^x e^y$ を満たします（<Ref to="mathematics/hs-analysis/elementary-functions#thm-exp-properties" /> の 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$。

</Proof>

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 以上であることをコーシー・シュワルツの不等式で示します。

<Proposition id="prop-softmax-jacobian" title="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$。

</Proposition>

<Proof of="prop-softmax-jacobian">

$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$ にコーシー・シュワルツの不等式（<Ref to="mathematics/linear-algebra/inner-product-spaces#thm-cauchy-schwarz" />）を使うと $\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$（<Ref to="prop-softmax-basic" /> の 1 で $\sum_i\sigma_i = 1$）。よって $\boldsymbol{v}^{\mathsf{T}}J\boldsymbol{v} \ge 0$ で、半正定値です（<Ref to="mathematics/linear-algebra/spectral-theorem#def-positive-definite" />）。

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^{*}}$ です。

</Proof>

2 は平行移動不変性（<Ref to="prop-softmax-basic" /> の 2）の微分形です。$\boldsymbol{1}$ の向きに動かしても出力が変わらないので、その向きの微分は 0 になります。3 は、softmax の出力がほぼ 1 点に集中している（$\max_k \sigma_k$ が 1 に近い）と、ヤコビ行列の全成分が 0 に近くなることを言っています。そうなると softmax を通る勾配がほとんど流れません。交差エントロピー（予測した確率が正解からどれだけ外れているかを測る損失）と組み合わせたときの勾配 $\boldsymbol{\sigma} - \boldsymbol{y}$（<Ref to="computer-science/math-for-ml/backpropagation#prop-softmax-ce" />）は、このヤコビ行列に損失の勾配を掛けた結果です（ヤコビ行列の式 $\mathrm{diag}(\boldsymbol{\sigma}) - \boldsymbol{\sigma}\boldsymbol{\sigma}^{\mathsf{T}}$ そのものは同じ命題にも述べられています。ここで新しく示したのは 1〜3 の性質です）。

<Proposition id="prop-mask-as-limit" title="マスクは極限として正当化できる">

$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}$ である。

</Proposition>

<Proof of="prop-mask-as-limit">

$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$。<Ref to="mathematics/hs-analysis/elementary-functions#thm-exp-properties" /> の 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 に収束します。正であることは <Ref to="prop-softmax-basic" /> の 1 です。

</Proof>

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

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

<Definition id="def-attention" title="スケール付き内積注意・自己注意・マルチヘッド・因果マスク">

**（注意）** $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$）である。

</Definition>

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

<Proposition id="prop-attention-convex-combination" title="注意の出力は凸結合">

<Ref to="def-attention" /> の注意（因果マスクつきの場合を含む）について、出力の第 $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$ である。

</Proposition>

<Proof of="prop-attention-convex-combination">

マスクなしのとき、$A$ の第 $i$ 行は $\sigma(\boldsymbol{s}_i)$ で、<Ref to="prop-softmax-basic" /> の 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$ です。

ノルムについて、三角不等式（<Ref to="mathematics/linear-algebra/inner-product-spaces#cor-triangle" />）と $\|c\boldsymbol{v}\| = |c|\|\boldsymbol{v}\|$（<Ref to="mathematics/linear-algebra/inner-product-spaces#def-norm" />）より

$$
\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$。上からも同様です。

</Proof>

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

<Example id="ex-attention-by-hand" title="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 つの行ベクトルが作る三角形の内側の点です（<Ref to="prop-attention-convex-combination" />）。

**行の入れ替え。** 第 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 節の因果性です。

</Example>

### 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$ も独立」といった確認は、その足し算を使ってよいことの裏づけなので、初めは結論だけ追っても構いません。

<Proposition id="prop-scaling-variance" title="内積の分散は次元に比例する">

$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
$$

が成り立つ。

</Proposition>

<Proof of="prop-scaling-variance">

独立な族の一部を取り出したものも独立です（<Ref to="mathematics/probability/random-variables#def-independence" /> で、取り出さない変数に対応するボレル集合を $\mathbb{R}$ に取る）。分散が有限なので各変数は 2 乗可積分で、したがって可積分です。

**平均。** $q_i$ と $k_i$ は独立で可積分なので、<Ref to="mathematics/probability/random-variables#prop-independence-product" /> より $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$（<Ref to="mathematics/probability/random-variables#prop-variance-properties" /> の 1）で、$k_i$ も同様なので、再び <Ref to="mathematics/probability/random-variables#prop-independence-product" /> より $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 個の可積分な確率変数なので、<Ref to="lem-product-expectation" />（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$。

**合計。** <Ref to="mathematics/probability/random-variables#prop-variance-properties" /> の 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$ です。

</Proof>

<Example id="ex-scaling-numeric" title="割らないと重みが 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 |

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

なお、ここでの「成分が独立で分散 1」は初期化の直後を想定した模型で、学習が進んだあとのクエリとキーがこの仮定を満たすわけではありません。単位球面上の一様な方向どうしの内積の分散が $1/d$ になること（<Ref to="computer-science/ai-era/relearning-mathematics#prop-inner-product-variance" />）は、長さを 1 にそろえた場合の同じ現象です。

</Example>

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

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

<Lemma id="lem-equivariant-closure" title="置換同変な写像の組み立て">

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)$ も置換同変である。

</Lemma>

<Proof of="lem-equivariant-closure">

$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)]$。最後の主張は、恒等写像が置換同変であることと和の場合から従います。

</Proof>

<Theorem id="thm-attention-equivariance" title="自己注意の置換同変性">

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

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

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

</Theorem>

<Proof of="thm-attention-equivariance">

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

**クエリ・キー・バリュー。** 積の結合法則（<Ref to="mathematics/linear-algebra/matrices-and-linear-systems#cor-associativity" />）と $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 の置換同変性（<Ref to="prop-softmax-basic" /> の 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}}$ を合成したものなので、<Ref to="lem-equivariant-closure" /> より置換同変です。後半の主張も同じ補題から、組み立ての段数についての帰納法で従います。

</Proof>

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

<Corollary id="cor-bag-of-tokens" title="位置情報がなければ語順を区別できない">

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

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

</Corollary>

<Proof of="cor-bag-of-tokens">

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$ とした場合です。

</Proof>

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

<Remark id="rem-causal-not-equivariant" title="因果マスクつきの場合">

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

</Remark>

## 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 行目（最後の行）と等しい、ということです。

<Definition id="def-causal-map" title="因果的な写像の族">

各 $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$ 行からなる行列である。

</Definition>

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

<Lemma id="lem-causal-closure" title="因果的な写像の組み立て">

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})$。

</Lemma>

<Proof of="lem-causal-closure">

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)$。両者は等しい。

3. $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$。よって両者は等しい。

2. 合成について。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$ が恒等写像）で因果的であることと和の場合からです。

</Proof>

<Theorem id="thm-causal-mask" title="因果マスクつきの自己注意は未来に依存しない">

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

</Theorem>

<Proof of="thm-causal-mask">

**1 ヘッド。** 長さ $n$ の入力 $X$ に対する出力の第 $i$ 行は、<Ref to="def-attention" /> より

$$
\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}}$ との合成も因果的です（<Ref to="lem-causal-closure" /> の 1 と 2）。

**層の重ね合わせ。** 層の数 $L$ についての帰納法です。$L = 1$ は上で示した場合か、<Ref to="lem-causal-closure" /> の 1 の形です。$L - 1$ 層までの合成が因果的なら、それに因果的な層を合成・加算したものも <Ref to="lem-causal-closure" /> の 2 より因果的です。

</Proof>

<Corollary id="cor-parallel-next-token" title="1 回の順伝播で n 個の予測が得られる">

$f = (f_n)$ を因果的な族とする（<Ref to="thm-causal-mask" /> の形のモデル）。長さ $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 回の順伝播で同時に、別々に計算したときと同じ値で得られる。

</Corollary>

<Proof of="cor-parallel-next-token">

<Ref to="def-causal-map" /> の定義そのものです。モデルが因果的であることは <Ref to="thm-causal-mask" /> です。

</Proof>

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

<Figure caption="n = 4 の因果マスク。行が「見る側」（クエリ）の位置 i、列が「見られる側」（キー）の位置 j。塗った成分（j ≤ i）だけ重みが正で、ほかは 0。重みの行列は下三角になる。">
<svg viewBox="0 0 520 250" width="100%" role="img" aria-label="4 行 4 列の因果マスク。下三角の成分だけが塗られている">
  <g fill="var(--sl-color-accent)" opacity="0.35">
    <rect x="180" y="40" width="44" height="44" />
    <rect x="180" y="84" width="88" height="44" />
    <rect x="180" y="128" width="132" height="44" />
    <rect x="180" y="172" width="176" height="44" />
  </g>
  <g fill="none" stroke="currentColor" stroke-width="1.2">
    <rect x="180" y="40" width="176" height="176" />
    <line x1="224" y1="40" x2="224" y2="216" />
    <line x1="268" y1="40" x2="268" y2="216" />
    <line x1="312" y1="40" x2="312" y2="216" />
    <line x1="180" y1="84" x2="356" y2="84" />
    <line x1="180" y1="128" x2="356" y2="128" />
    <line x1="180" y1="172" x2="356" y2="172" />
  </g>
  <g fill="currentColor" font-size="14" text-anchor="middle">
    <text x="202" y="30">j = 1</text>
    <text x="246" y="30">2</text>
    <text x="290" y="30">3</text>
    <text x="334" y="30">4</text>
    <text x="150" y="67">i = 1</text>
    <text x="150" y="111">2</text>
    <text x="150" y="155">3</text>
    <text x="150" y="199">4</text>
    <text x="246" y="67">0</text>
    <text x="290" y="67">0</text>
    <text x="334" y="67">0</text>
    <text x="290" y="111">0</text>
    <text x="334" y="111">0</text>
    <text x="334" y="155">0</text>
    <text x="440" y="111">−∞ を足す</text>
    <text x="440" y="131">（重み 0）</text>
    <text x="80" y="155">重みが正</text>
  </g>
</svg>
</Figure>

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

<Proposition id="prop-attention-parameters" title="マルチヘッド自己注意のパラメータ数">

<Ref to="def-attention" /> のマルチヘッド自己注意（$d_k = d_v = d/h$、バイアスあり）1 層のパラメータ数は

$$
4d^2 + 4d
$$

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

</Proposition>

<Proof of="prop-attention-parameters">

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$ です。

</Proof>

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

<Proposition id="prop-attention-complexity" title="自己注意 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 倍になる。

</Proposition>

<Proof of="prop-attention-complexity">

$p \times q$ 行列と $q \times r$ 行列の積の乗算の回数は $pqr$ です（<Ref to="computer-science/research-computing/numpy-array-computing#def-einsum" />）。

- 射影 $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)$（<Ref to="computer-science/algorithms/complexity-and-big-o#def-big-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$ です（<Ref to="computer-science/algorithms/complexity-and-big-o#prop-doubling" /> の 3）。

</Proof>

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 つの項が等しくなります。

<Remark id="rem-cross-attention" title="交差注意">

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

</Remark>

## 8. 最小実装と数値検証

マルチヘッド自己注意を NumPy で実装します。ヘッドの軸を先頭に持つ 3 次元配列でパラメータを持ち、積は `einsum`（<Ref to="computer-science/research-computing/numpy-array-computing#def-einsum" />）で書きます。バイアスの加算と $\sqrt{d_k}$ での除算はブロードキャスト（<Ref to="computer-science/research-computing/numpy-array-computing#def-broadcast" />）です。

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

```python
# 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()))
```

```text
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 行目：パラメータの総数が <Ref to="prop-attention-parameters" /> の $4d^2 + 4d = 288$ に一致。
- 2 行目：マスクなしでは $f(PX) = Pf(X)$ が丸め誤差の範囲で成り立つ（<Ref to="thm-attention-equivariance" />）。
- 3 行目：因果マスクつきでは成り立たない（<Ref to="rem-causal-not-equivariant" />）。
- 4・5 行目：5・6 番目の入力を変えても、因果マスクつきなら 1〜4 番目の出力はまったく変わらず、マスクなしでは変わる（<Ref to="thm-causal-mask" />）。
- 6 行目：接頭辞だけを入力したときの最後の行が、全体を入力したときの対応する行に一致（<Ref to="cor-parallel-next-token" />）。
- 7・8 行目：重みが非負で行の和が 1、上三角が 0、ヘッドの出力の行のノルムがバリューの行のノルムの最大値以下（<Ref to="prop-attention-convex-combination" />）。

<Ref to="ex-scaling-numeric" /> の表と 7 節の回数の表は、次のコードの出力です。

```python
# 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)
```

```text
  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 ブロック](/computer-science/llm-architecture/transformer-block)では、自己注意に残差接続（入力に層の出力を足す接続）・[層正規化](/computer-science/llm-architecture/transformer-block#def-layer-norm)・行ごとの MLP を組み合わせて 1 つの層にし、それを重ねてモデル全体を作ります。<Ref to="lem-equivariant-closure" /> と <Ref to="lem-causal-closure" /> は「行ごとに働く写像との合成と和で性質が保たれる」ことを言っているので、同変性と因果性はブロック全体にそのまま引き継がれます。

## 10. 演習

<Exercise id="exr-softmax-compute" difficulty="易">

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

<Solution>

$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}$ なので、平行移動不変性（<Ref to="prop-softmax-basic" /> の 2）より $\sigma(\boldsymbol{z}') = (1/6, 1/3, 1/2)$ です。定義どおりに $e^{100}$ を計算する必要はありません。

</Solution>
</Exercise>

<Exercise id="exr-jacobian-two" difficulty="標準">

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

<Solution>

<Ref to="prop-softmax-jacobian" /> より
$$
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 に近づきます。出力が一方に偏るほど勾配が流れにくくなることが、この場合は式で見えます。

</Solution>
</Exercise>

<Exercise id="exr-mask-limit" difficulty="易">

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

<Solution>

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

</Solution>
</Exercise>

<Exercise id="exr-causal-last-row" difficulty="標準">

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

<Solution>

<Ref to="thm-causal-mask" /> の証明の式で $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
# 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())
```

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

</Solution>
</Exercise>

## 参考文献

- [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](https://arxiv.org/abs/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](https://arxiv.org/abs/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](https://arxiv.org/abs/1810.00825) — 注意に基づく置換同変・置換不変なネットワークの定式化。
- [4] 本サイト「[確率変数と期待値：可測関数としての確率変数、ルベーグ積分としての期待値](/mathematics/probability/random-variables)」 — 独立性、分散の性質、独立な確率変数の積の期待値。

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

**<Ref to="prop-scaling-variance" /> の証明で使った、3 個以上の積についての補題を示します。** 2 個の場合は既出の命題で、ここではそれを個数についての帰納法で延長します。

<Lemma id="lem-product-expectation" title="独立な可積分確率変数の積の期待値（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]
$$

が成り立つ。

</Lemma>

<Proof of="lem-product-expectation">

$m$ についての帰納法です。$m = 2$ は <Ref to="mathematics/probability/random-variables#prop-independence-product" /> です。$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$ が確率変数であることは <Ref to="mathematics/probability/random-variables#prop-operations" /> です。

**$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})$ を生成するので（下で使うのと同じ <Ref to="mathematics/real-analysis/product-measures-fubini#prop-borel-product" />）、すべてのボレル集合がこの性質を持ちます。$\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})$ を生成します（<Ref to="mathematics/real-analysis/product-measures-fubini#prop-borel-product" />）。長方形の上では、$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)
$$

です。測度の一意性（<Ref to="mathematics/real-analysis/abstract-measures#thm-uniqueness-measures" />）より、$\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$ は独立で可積分なので、<Ref to="mathematics/probability/random-variables#prop-independence-product" /> より $X_1Y$ は可積分で $E[X_1Y] = E[X_1]E[Y] = E[X_1]E[X_2]\cdots E[X_m]$ です。

</Proof>