セルフアテンションとクロスアテンション
導入
名前とは異なり、セルフアテンションself-attentionとクロスアテンションcross-attentionはアテンション関数それ自体についての定義ではない。アテンション関数とは、クエリ行列$\mathbf{Q}$、キー行列$\mathbf{K}$、バリュー行列$\mathbf{V}$が与えられているとき、次のように計算される関数のことである。
$$ \operatorname{Attention}(\mathbf{Q},\mathbf{K},\mathbf{V}) = \mathbf{V} \operatorname{Softmax}\left(f(\mathbf{Q},\mathbf{K})\right) \tag{1} $$
- $\mathbf{Q} \in M^{d_{k} \times n}$
- $\mathbf{K} \in M^{d_{k} \times m}$
- $\mathbf{V} \in M^{d_{v} \times m}$
このとき$\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} $$
$(1)$で$f$を具体的にどう定義するかによってアテンション関数が変わるが、セルフアテンションとクロスアテンションはアテンション関数を定義する方法ではなく、クエリ/キー/バリュー行列を定義する方法に関するものである。
定義
データ行列$\mathbf{X} \in M^{d \times n}$と、学習される重み行列$\mathbf{W}^{Q} \in M^{d_{k} \times d}$、$\mathbf{W}^{K} \in M^{d_{k} \times d}$、$\mathbf{W}^{V} \in M^{d_{v} \times d}$が与えられたとしよう。クエリ、キー、バリュー行列をすべて同じデータ$\mathbf{X}$から得る方法をセルフアテンションという。
$$ \text{Self-Attention:} \quad \mathbf{Q} = \mathbf{W}^{Q}\mathbf{X}, \quad \mathbf{K} = \mathbf{W}^{K}\mathbf{X}, \quad \mathbf{V} = \mathbf{W}^{V}\mathbf{X} $$
一方、互いに異なる二つのデータ行列$\mathbf{X} \in M^{d_{x} \times n}$、$\mathbf{Y} \in M^{d_{y} \times m}$に対して、クエリ行列は$\mathbf{X}$から、キーとバリュー行列は$\mathbf{Y}$から得る方法をクロスアテンションという。
$$ \text{Cross-Attention:} \quad \mathbf{Q} = \mathbf{W}^{Q}\mathbf{X}, \quad \mathbf{K} = \mathbf{W}^{K}\mathbf{Y}, \quad \mathbf{V} = \mathbf{W}^{V}\mathbf{Y} $$
このとき重み行列の大きさは$\mathbf{W}^{Q} \in M^{d_{k} \times d_{x}}$、$\mathbf{W}^{K} \in M^{d_{k} \times d_{y}}$、$\mathbf{W}^{V} \in M^{d_{v} \times d_{y}}$である。
説明
アテンション関数の形にまとめると以下のとおりである。
$$ \begin{align*} \text{Self-Attention:} \quad &\operatorname{Attention}(\mathbf{W}^{Q}\mathbf{X}, \mathbf{W}^{K}\mathbf{X}, \mathbf{W}^{V}\mathbf{X}) \\ &= \mathbf{W}^{V}\mathbf{X} \operatorname{Softmax}\left( \left( \mathbf{W}^{K}\mathbf{X} \right)^{\mathsf{T}} \left( \mathbf{W}^{Q}\mathbf{X} \right) \right) \end{align*} $$
$$ \begin{align*} \text{Cross-Attention:} \quad &\operatorname{Attention}(\mathbf{W}^{Q}\mathbf{X}, \mathbf{W}^{K}\mathbf{Y}, \mathbf{W}^{V}\mathbf{Y}) \\ &= \mathbf{W}^{V}{\color{red}\mathbf{Y}} \operatorname{Softmax}\left( \left( \mathbf{W}^{K}{\color{red}\mathbf{Y}} \right)^{\mathsf{T}} \left( \mathbf{W}^{Q}\mathbf{X} \right) \right) \end{align*} $$
セルフアテンションは、シーケンスが自分自身にアテンションを取ることである。クエリ、キー、バリューがすべて同じデータ$\mathbf{X}$から出てくるので、クエリの個数とキー/バリューの個数が同じであり($m = n$)、$\mathbf{X}$の各列ベクトル、例えば文の各単語が同じ文の中の他のすべての単語とスコアを測り、その比重どおりにバリューを混ぜる。その結果、各単語の表現が自分の属する文全体の文脈を反映するように更新される。たとえば「それ」のような代名詞が文の中のどの名詞を指しているかが、セルフアテンションの重みとして現れる。トランスフォーマーのエンコーダが代表的にセルフアテンションを使用しており、イントラアテンションintra-attentionとも呼ばれる。
クロスアテンションは、互いに異なる二つのシーケンスの間で一方が他方を参照することである。機械翻訳が代表的な例で、デコーダが作っている翻訳文側の特徴ベクトルがクエリ$\mathbf{X}$になり、エンコーダが出した原文側の特徴ベクトルがキーとバリュー$\mathbf{Y}$になる。再帰型ニューラルネットワークに基づくseq2seqモデルに初めてアテンションを導入したバダナウ・アテンションがまさにこの形であり、トランスフォーマーのデコーダでエンコーダの出力を参照するアテンション層もクロスアテンションである。このような文脈からクロスアテンションはエンコーダ・デコーダアテンションencoder-decoder attentionとも呼ばれる。