Attentionと、全結合・畳み込みとの関係
Transformerに使われるAttentionと、FFN(Position-wise Dense)、全結合(Dense)、畳み込み(Convolution)の 関係を俯瞰するために、お絵描きを試みたので、ここに載せる。
もしかしたら不正確な表現があるかもしれないが、ご容赦いただきたい。
なお、図では、下側を入力、上側を出力としている。
(2021/6/30: 図のインデックスの誤りを修正)
1. Dense (全結合)
全結合は、基本的に重み行列$ \bm W $の乗算で表される。
なお、通常はもうひとつ別の重みベクトル(バイアス)$\bm b$を用意して、 $\bm y = \bm W \bm x + \bm b$とするが、本記事では省略する。
2. Convolution (畳み込み)
簡単のため、1次元の畳み込み演算(カーネルサイズ3、入出力の次元は同じ)とする。
畳み込み演算は、一部の重み(1次元畳み込みの場合は、行列の対角方向の要素)が共有され、かつ、 重みの外側がゼロである行列を用意することで、全結合と同様に行列演算として表すことができる。
バイアスを組み入れることの方が多いが、その場合は$\bm y = \bm C \bm x + \bm b$とする。
3. Position-wise Dense
ここから2次元(Depth軸とPosition軸)。
Transformerブロックでも使われるPosition-wise Denseでは、 Depth軸だけを入力とした全結合(重み$\tilde \bm W$)を、 全てのPosition位置で行う。このとき、重み$\tilde \bm W$は、Position位置に依らず共有とする。
4. Depth-wise Convolution
Transformerとは関係ないが、比較対象として。
Depth-wise Convolutionでは、全てのDepth位置で、Position方向に広がった畳み込みを行う。
畳み込む対象がスカラーからベクトルになった、と考えた方が分かりやすいかも。
なお、Depth-wise Convolutionの2次元版の応用例としては、 フルのConvolutionをDepth-wiseとPoint-wiseに分割しすることによる パラメータ数削減(モデル軽量化)が挙げられる。
5. Self Attention / Source-Target Attention
以上を踏まえて、Attentionの基本的な部分を図にすると、このような感じ。
ValueもKeyもQueryも、Position-wise Denseに相当する方法で生成する。
Depth-wise Convolutionのような、ベクトルを重み付けして集約する機能があり、 その重み付けは、QueryとKeyの内積ベースで決まる。
完全に全結合として行数×列数分の要素を重みパラメータとするよりも、 行(q)と列(k)に分けて用意する方が、(少なくとも)重みパラメータは少なくて済むことが多いはず。
6. Multi-headed Attention
ついでに、Multi-headの場合。
図では全てを示せていないが、Valueだけではなく、KeyとQueryも同じように分割した後、 Attentionの計算をして、得られる結果を単純に結合する。