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

（注意） $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$ は行ごと）。$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}$ に対し

として $\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}$ を使って

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

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