About this article
tf.keras.layers.MultiHeadAttentionの理解を深めるための勉強メモです.
1. Understanding Transformer
- まず,基本的な関数の引数の数式の次元から理解をする.
- 次に,その関数に実際何が代入されるかを理解する.
1.a. Looking at expression
上記の論文をみると,AttentionとMulti-Head-Attentionの条件がわかる.
1.a. As to Attention especially dimension
Attention(A)の引数は三つであり以下の条件があります:
-
$V,K,Q$は行列である.
- それゆえ,それぞれ $Q\in\mathbb{R}^{q_1\times q_2}$ のように書くことにする.
-
(1)の式 より $QK^{\rm T}, K^{\rm T}V$から $q_2 = k_2=:A,v_1=k_1=:B$.
-
よって, $V\in\mathbb{R}^{B\times v_2}, K\in\mathbb{R}^{B\times A},Q\in\mathbb{R}^{q_1\times A}$
- Q(uery)の$q_1$やV(alue)の$v_2$はinputやoutputから決まる.
- ただし,(1)の数式から出力は$\mathbb{R}^{q_1 \times v_2}$ゆえ,$A,B$はモデルを作成する人が自由に設定できるハイパーパラメータ.
1.b. As to Attention especially dimension
Multi-Head-Attention(MHA)の引数は三つであり以下の条件があります:
- $V,K,Q$は行列である.
- p.5の式 より $v_2 = k_2 = q_2 = d_{\rm model}, v_1 = k_1 := C$.
- よって,$V\in\mathbb{R}^{C\times d_{\rm model}}, K\in\mathbb{R}^{C\times d_{\rm model}},Q\in\mathbb{R}^{q_1\times d_{\rm model}}$
- また,出力は$\mathbb{R}^{q_1 \times q_2}$ (ここでは,$d_{\rm{model}}=q_2$であるが,3章で解説する)
- Attentionとしては,Headの数だけ$\mathbb{R}^{q_1 \times d_v }$ができる.
2. Looking at Architectures
それぞれのMHA(Multi-Head-Attention)に注目します.
-
p.3 によると EncoderのMHAはSelf-attention machine とあります.
- つまり,1つ目のencoderに入るものとしては,embeddingされた入力値を$X\in\mathbb{R}^{l_{\rm in}\times d_{\rm in}}$とすると,$X=V=K=Q$であり,$C=q_1=l_{\rm in},d_{\rm model}=d_{\rm in}$ である.
- また,出力は$\mathbb{R}^{l_{\rm in}\times d_{\rm in}}$
-
p.3 によると それぞれのDecoderの一つ目のMHAはMasked-MHAである.
- 目的の出力を$Y\in\mathbb{R}^{l_{\rm out}\times d_{\rm out}}$とすると,$i\in[l]$を予測する時は,$l$未満を知っている状態にする必要があり,p.5によると$l$以上の値は$-\infty$で値を調節して,自己回帰的にdecodeを行う.
- また,ひとつ目のDecoderにおいては$Y=V=K=Q$で$C=q_1=l_{\rm out},d_{\rm model}=d_{\rm out}$ である.
- また,出力は$\mathbb{R}^{l_{\rm out}\times d_{\rm out}}$
-
p.3 によると それぞれのDecoderの二つ目のMHAは通常のものである.
- 図によると,Encoderの出力を$Z\in \mathbb{R}^{l_{\rm in}\times d_{\rm in}}$,Decoderのひとつ目の出力を$Y'\in \mathbb{R}^{l_{\rm out}\times d_{\rm out}}$とすると, $V=K=Z, Q=Y'$であるため,1.b.より $d_{\rm model}=d_{\rm in}=d_{\rm out}$である必要がある.
- また,$C=l_{\rm in}, q_1 = l_{\rm out}$である.
- また,出力は$\mathbb{R}^{l_{\rm out}\times d_{\rm model}}$
- 図によると,Encoderの出力を$Z\in \mathbb{R}^{l_{\rm in}\times d_{\rm in}}$,Decoderのひとつ目の出力を$Y'\in \mathbb{R}^{l_{\rm out}\times d_{\rm out}}$とすると, $V=K=Z, Q=Y'$であるため,1.b.より $d_{\rm model}=d_{\rm in}=d_{\rm out}$である必要がある.
以上から本論文の手法を用いるためには,1.の条件に加えて入力(target)と出力(source)のembeddingの次元($d_{\rm in},d_{\rm out}$)を同じ値にする必要があることがわかりました.
3. Question
2.の結論を得ましたが,実際に$d_{\rm in},d_{\rm out}$を同じにしないといけないわけではないことがp.5. でわかります.
というのも(これは私の見解ですが),p.5.のパラメータを調節することで,こちらの統一をする必要がないためです.すなわち, 1.b. の条件は過剰であり わかりやすくするために用意されたものだと考えられるからです.
4. tf.keras.layers.MultiHeadAttention
HPに幾つか例が記載されているため,それを編集してみましょう.
ここでは,あえて$d_{\rm in}$と$d_{\rm out}$を異なる数字にして理解を深めることにします.
4. Example 1.
- ここでは,
d_in,d_outの次元をづらた際にd_outが出力にきていることを確認した.
>>> d_in, d_out = 7, 16
>>> l_in, l_out = 100, 121
>>> source = tf.keras.Input(shape=[l_in, d_in])
>>> target = tf.keras.Input(shape=[l_out, d_out])
>>> layer = MultiHeadAttention(num_heads=500, key_dim=2)
>>> output_tensor, weights = layer(query=target, value=source, return_attention_scores=True)
>>> print(output_tensor.shape, weights.shape)
(None, 121, 16) (None, 500, 121, 100)
- ここで,
sourceはK(ey),V(alue)でtargetはQ(uery)である.- このことから,このライブラリを使用する際には,embeddingの次元を無理に同じにする必要がないことが考えられる.
- また,ここで得られるweight(
attention_scores)は(num_heads,l_out,l_in)からp.5の式から得られるattentionを計算する前の,valueをかける前の$(QW_i^{Q})(KW_i^{K})^{\rm T}$をsoftmax等でスケールしたものである(と考えられる).
4. Example 2.
HPにてattention_axesの使用例がありました.
これを行うことで,attentionを限定的に行えるようです(仕様がよくわからなかったため詳細は割愛いたします).