tf.keras.layers.MultiHeadAttentionについて

# About this article

[tf.keras.layers.MultiHeadAttention](https://www.tensorflow.org/api_docs/python/tf/keras/layers/MultiHeadAttention)の理解を深めるための勉強メモです.

# 1. Understanding `Transformer`

- まず,基本的な関数の引数の数式の次元から理解をする.
- 次に,その関数に実際何が代入されるかを理解する.

https://arxiv.org/abs/1706.03762

# 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}}$

以上から本論文の手法を用いるためには,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](https://www.tensorflow.org/api_docs/python/tf/keras/layers/MultiHeadAttention)

[HP](https://www.tensorflow.org/api_docs/python/tf/keras/layers/MultiHeadAttention)に幾つか例が記載されているため,それを編集してみましょう.
ここでは,あえて$d_{\rm in}$と$d_{\rm out}$を異なる数字にして理解を深めることにします.

## 4. Example 1.

- ここでは,`d_in,d_out`の次元をづらた際に`d_out`が出力にきていることを確認した.
```python3
>>> 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等でスケールしたものである(と考えられる).
    - https://github.com/keras-team/keras/blob/v2.12.0/keras/layers/attention/multi_head_attention.py#L527-L531 
    - その他, keyがどのように設計されるなどは こちらのコードを参照のこと : https://github.com/keras-team/keras/blob/v2.12.0/keras/layers/attention/multi_head_attention.py#L130-L730


## 4. Example 2.
[HP](https://www.tensorflow.org/api_docs/python/tf/keras/layers/MultiHeadAttention)にて`attention_axes`の使用例がありました.

これを行うことで,attentionを限定的に行えるようです(仕様がよくわからなかったため詳細は割愛いたします).

<!--

```python3
>>> 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, attention_axes=0)
>>> output_tensor, weights = layer(query=target, value=source, return_attention_scores=True)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "**/.venv/lib/python3.9/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "**/.venv/lib/python3.9/site-packages/tensorflow/python/framework/ops.py", line 1973, in _create_c_op
    raise ValueError(e.message)
ValueError: Exception encountered when calling layer "multi_head_attention_17" (type MultiHeadAttention).

Dimensions must be equal, but are 100 and 121 for '{{node multi_head_attention_17/einsum/Einsum}} = Einsum[N=2, T=DT_FLOAT, equation="ebcd,abcd->bcae"](multi_head_attention_17/key/add, multi_head_attention_17/Mul)' with input shapes: [?,100,500,2], [?,121,500,2].

Call arguments received by layer "multi_head_attention_17" (type MultiHeadAttention):
  • query=tf.Tensor(shape=(None, 121, 16), dtype=float32)
  • value=tf.Tensor(shape=(None, 100, 7), dtype=float32)
  • key=None
  • attention_mask=None
  • return_attention_scores=True
  • training=None
  • use_causal_mask=False
```

```python3
>>> 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, attention_axes=[1])
>>> 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)
```


```python3
>>> 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, attention_axes=[2])
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "**/.venv/lib/python3.9/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "**/.venv/lib/python3.9/site-packages/tensorflow/python/framework/ops.py", line 1973, in _create_c_op
    raise ValueError(e.message)
ValueError: Exception encountered when calling layer "multi_head_attention_21" (type MultiHeadAttention).

Dimensions must be equal, but are 100 and 121 for '{{node multi_head_attention_21/einsum/Einsum}} = Einsum[N=2, T=DT_FLOAT, equation="abed,abcd->abce"](multi_head_attention_21/key/add, multi_head_attention_21/Mul)' with input shapes: [?,100,500,2], [?,121,500,2].

Call arguments received by layer "multi_head_attention_21" (type MultiHeadAttention):
  • query=tf.Tensor(shape=(None, 121, 16), dtype=float32)
  • value=tf.Tensor(shape=(None, 100, 7), dtype=float32)
  • key=None
  • attention_mask=None
  • return_attention_scores=True
  • training=None
  • use_causal_mask=False
```


```python3
>>> 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, attention_axes=[3])
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "**/.venv/lib/python3.9/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "**/.venv/lib/python3.9/site-packages/tensorflow/python/ops/special_math_ops.py", line 1287, in _einsum_v2_parse_and_resolve_equation
    raise ValueError(
ValueError: Exception encountered when calling layer "multi_head_attention_22" (type MultiHeadAttention).

Output subscripts contain a label appearing more than once: abcd,abcd->abcdd

Call arguments received by layer "multi_head_attention_22" (type MultiHeadAttention):
  • query=tf.Tensor(shape=(None, 121, 16), dtype=float32)
  • value=tf.Tensor(shape=(None, 100, 7), dtype=float32)
  • key=None
  • attention_mask=None
  • return_attention_scores=True
  • training=None
  • use_causal_mask=False
```


-->