スケールド・ドットプロダクト・アテンション
導入
それぞれクエリ・キー・バリューと呼ばれる三つの行列$\mathbf{Q} \in M^{d_{k} \times n}$、$\mathbf{K} \in M^{d_{k} \times m}$、$\mathbf{V} \in M^{d_{v} \times m}$と、クエリとキーに対するスコア関数$f : M^{d_{k} \times n} \times M^{d_{k} \times m} \to M^{m \times n}$が与えられたとしよう。アテンションとは、次のように定義される関数をいう。
$$ \begin{align*} \operatorname{Attention}: M^{d_{k} \times n} \times M^{d_{k} \times m} \times M^{d_{v} \times m} &\to M^{d_{v} \times n} \\ (\mathbf{Q},\mathbf{K},\mathbf{V}) &\mapsto \mathbf{V} \operatorname{Softmax}\left(f(\mathbf{Q},\mathbf{K})\right) \tag{1} \end{align*} $$
ここで$\operatorname{Softmax}$とは、与えられた行列$\mathbf{X} = \begin{bmatrix} \mathbf{x}_{1} & \cdots & \mathbf{x}_{N}\end{bmatrix}$に対して各列ベクトルにソフトマックス $\operatorname{softmax}$を取り、各列の和が$1$になるようにする関数をいう。
$$ \operatorname{Softmax}(\mathbf{X}) := \begin{bmatrix} \underset{\vert}{\overset{\vert}{\operatorname{softmax}(\mathbf{x}_{1})}} & \cdots & \underset{\vert}{\overset{\vert}{\operatorname{softmax}(\mathbf{x}_{N})}} \end{bmatrix} $$
$f$をどのように定義するかによってアテンションの種類が分かれるが、Ashish Vaswaniほか7名がトランスフォーマー構造を最初に提案した論文『Attention Is All You Need』1で提案された形のアテンションは以下の通りである。
定義
以下の形のアテンション関数をスケールド・ドットプロダクト・アテンションscaled dot-product attentionという。
$$ \operatorname{Attention} (\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \mathbf{V} \operatorname{Softmax} \left( \frac{\mathbf{K}^{\mathsf{T}} \mathbf{Q}}{\sqrt{d_{k}}} \right) $$
説明
名前の通り、クエリベクトル$\mathbf{q}$とキーベクトル$\mathbf{k}$のスコア関数を、スケーリングファクターを含む内積としたものである。
$$ \operatorname{score}(\mathbf{q}, \mathbf{k}) = \dfrac{1}{\sqrt{d_{k}}} \mathbf{k}^{\mathsf{T}} \mathbf{q} $$
$$ f(\mathbf{Q}, \mathbf{K}) = \frac{\mathbf{K}^{\mathsf{T}} \mathbf{Q}}{\sqrt{d_{k}}} = \frac{1}{\sqrt{d_{k}}} \begin{bmatrix} \mathbf{k}_{1} \cdot \mathbf{q}_{1} & \mathbf{k}_{1} \cdot \mathbf{q}_{2} & \cdots & \mathbf{k}_{1} \cdot \mathbf{q}_{n} \\ \mathbf{k}_{2} \cdot \mathbf{q}_{1} & \mathbf{k}_{2} \cdot \mathbf{q}_{2} & \cdots & \mathbf{k}_{2} \cdot \mathbf{q}_{n} \\ \vdots & \vdots & \ddots & \vdots \\ \mathbf{k}_{m} \cdot \mathbf{q}_{1} & \mathbf{k}_{m} \cdot \mathbf{q}_{2} & \cdots & \mathbf{k}_{m} \cdot \mathbf{q}_{n} \end{bmatrix} $$
なぜ$\sqrt{d_{k}}$で割るのか
内積をそのままスコアとして使わず、わざわざ$\sqrt{d_{k}}$で割る理由は、次元$d_{k}$が大きくなるほど内積の大きさが$\sqrt{d_{k}}$のスケールで大きくなり、ソフトマックスを飽和させるからである1。クエリ$\mathbf{q}$とキー$\mathbf{k}$の各成分$q_{i}$、$k_{i}$がすべて平均$0$、分散$1$であり、互いに独立な確率変数であると仮定しよう。すると期待値の線形性と独立性により、内積$\mathbf{k}^{\mathsf{T}} \mathbf{q} = \sum_{i=1}^{d_{k}} k_{i} q_{i}$の平均は次のように$0$である。
$$ \mathbb{E} \left[ \mathbf{k}^{\mathsf{T}} \mathbf{q} \right] = \mathbb{E} \left[ \sum_{i=1}^{d_{k}} k_{i} q_{i} \right] = \sum_{i=1}^{d_{k}} \mathbb{E} [ k_{i} q_{i} ] = \sum_{i=1}^{d_{k}} \mathbb{E} [ k_{i} ] \mathbb{E} [ q_{i} ] = 0 $$
成分がすべて互いに独立なので、積$k_{i} q_{i}$たちも互いに独立であり、独立な確率変数の和の分散は各分散の和である。また平均が$0$なので$\mathbb{E} [ k_{i}^{2} ] = \Var [ k_{i} ] = 1$であり、内積の分散は次のようにちょうど$d_{k}$になる。
$$ \begin{align*} \Var \left[ \mathbf{k}^{\mathsf{T}} \mathbf{q} \right] &= \sum_{i=1}^{d_{k}} \Var [ k_{i} q_{i} ] = \sum_{i=1}^{d_{k}} \left( \mathbb{E} [ k_{i}^{2} q_{i}^{2} ] - \left( \mathbb{E} [ k_{i} q_{i} ] \right)^{2} \right) \\ &= \sum_{i=1}^{d_{k}} \mathbb{E} [ k_{i}^{2} ] \mathbb{E} [ q_{i}^{2} ] = \sum_{i=1}^{d_{k}} \Var [ k_{i} ] \Var [ q_{i} ] = d_{k} \end{align*} $$
言い換えれば、スコア$\mathbf{k}^{\mathsf{T}} \mathbf{q}$の標準偏差が$\sqrt{d_{k}}$だということであり、次元が$d_{k} = 512$ならば標準偏差は$\sqrt{512} \approx 22.6$なので、絶対値が数十に達するスコアが出てくることになる。絶対値の大きい値がソフトマックスに入ると、指数関数が最大値の成分を圧倒的に大きくして出力がワンホットベクトルに近く飽和するが、飽和領域では入力がある程度変化しても出力がほとんど変わらないため、勾配が極めて小さくなり勾配消失によって学習が妨げられる。一方、スコアを$\sqrt{d_{k}}$で割れば分散が再び$1$に正規化されるので、$d_{k}$がいくら大きくなってもスコアの大きさは一定に保たれる。
$$ \Var \left[ \frac{\mathbf{k}^{\mathsf{T}} \mathbf{q}}{\sqrt{d_{k}}} \right] = \frac{1}{d_{k}} \Var \left[ \mathbf{k}^{\mathsf{T}} \mathbf{q} \right] = 1 $$
下の図がこの違いを示している。標準正規分布から成分をサンプリングして$d_{k} = 512$のクエリ$\mathbf{q}$を一つ、キー$\mathbf{k}_{1}, \dots, \mathbf{k}_{10}$を十個作り、スコア$\mathbf{k}_{j}^{\mathsf{T}} \mathbf{q}$十個をソフトマックスに通した結果である。左のようにスケーリングなしでそのまま通すと、確率が事実上キー一つに偏ってワンホットに飽和する。一方、同じスコアを$\sqrt{d_{k}} = \sqrt{512}$で割ってから通すと、右のように複数のキーに確率が広く分散し、アテンションが特定のキー一つだけを見るのではなく、複数のキーを幅広く参照できるようになる。
