[Attention入門]seq2seqとAttentionの解説(TensorFlow)

BERTを勉強していて**seq2seqのAttention**について整理しました。モデル部分は理解しましたが、**訓練ジョブを流す部分などはほとんど見ていないですし解説もしていません**。seq2seqについては[記事「【Keras入門(7)】単純なSeq2Seqモデル定義」](https://qiita.com/FukuharaYohei/items/27cd247342a0f7006511)で以前解説をしました。
もともとは以下のオンライン講座を受講していて、Attentionに関する理解を深めようとしたのがきっかけでした。

https://www.udemy.com/course/nlp-bert/

以下のTensorFlowのチュートリアルを使って学習しました。

https://www.tensorflow.org/text/tutorials/nmt_with_attention

チュートリアルに以下の記載があるとおり、少し**古い考えですが Attention を理解する上で重要**だと信じ時間をかけました。
> While this architecture is somewhat outdated it is still a very useful project to work through to get a deeper understanding of attention mechanisms (before going on to Transformers).

理解にあたり以下を参考にしました。私がPyTorchを使っていないことには、深い理由はないです。

https://qiita.com/halhorn/items/614f8fe1ec7663e04bea

https://qiita.com/m__k/items/646044788c5f94eadc8d

# モデル解説
## モデル全体
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/134273/e7119e18-5525-281d-99d5-82901780552c.png)

## モデル詳細
### Text Processing
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/134273/5427da40-8b18-f07e-af7e-005d2b2440c4.png)
テキストをTokenizeとID化しています(正確には以下の処理)。

1. ユニコード正規化
1. 大文字の小文字化
1. テキストのTokenize
1. [START]と[END]を先頭と末尾に追加
1. Tokenに対して辞書を使ったID化

出力のShapeに「最大Token数」とありますが、[START]と[END]が加わっていることが注意点です。「BS」はBatch Sizeを意味します。

|入出力|Shape|例|
|:-:|:-:|:-:|
|入力|BS * 1|[['I like you'],['This is a pen']]|
|出力|BS * 最大Token数|[[5 4 3 2 6 0],[5 1 1 1 1 6]]|

`TextVectorization`関数を使っています。確か日本語はTokenizeできないので、MeCabなどを使わないといけなかったはず。

```python
max_vocab_size = 5000

# Encoder
input_text_processor = tf.keras.layers.TextVectorization(
    standardize=tf_lower_and_split_punct,
    max_tokens=max_vocab_size)

# Decoder
output_text_processor = tf.keras.layers.TextVectorization(
    standardize=tf_lower_and_split_punct,
    max_tokens=max_vocab_size)
```

パラメータ`standardize`に渡した関数`tf_lower_and_split_punct`でユニコード正規化・小文字化などをしています。
ユニコード正規化では「㍍」を「メートル」に変換などをしています。軽くしか試していないので`tf_text.normalize_utf8`が日本語処理として実用に耐えうるレベルか不明。

```python:tf_lower_and_split_punct
def tf_lower_and_split_punct(text):
  # Split accecented characters.
  text = tf_text.normalize_utf8(text, 'NFKD')
  text = tf.strings.lower(text)
  # Keep space, a to z, and select punctuation.
  text = tf.strings.regex_replace(text, '[^ a-z.?!,¿]', '')
  # Add spaces around punctuation.
  text = tf.strings.regex_replace(text, '[.?!,¿]', r' \0 ')
  # Strip whitespace.
  text = tf.strings.strip(text)

  text = tf.strings.join(['[START]', text, '[END]'], separator=' ')
  return text
```

### Encoder
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/134273/7317f31b-2aa7-fd4a-7db2-5a3183638486.png)
Attention固有の内容は特になく通常のEncorderです。RNNセルにはGRUを使っています。

|入出力|Shape|
|:-:|:-:|
|Embedding入力|BS * 最大Token数|
|Embedding->RNN|BS * 最大Token数 * Embedding Unit数|
|RNN出力(State)|BS * RNN Unit数|
|RNN出力(Output)|BS * 最大Token数 * RNN Unit数|

```python:Encoder
class Encoder(tf.keras.layers.Layer):
  def __init__(self, input_vocab_size, embedding_dim, enc_units):
    super(Encoder, self).__init__()
    self.enc_units = enc_units
    self.input_vocab_size = input_vocab_size

    # The embedding layer converts tokens to vectors
    self.embedding = tf.keras.layers.Embedding(self.input_vocab_size,
                                               embedding_dim)

    # The GRU RNN layer processes those vectors sequentially.
    self.gru = tf.keras.layers.GRU(self.enc_units,
                                   # Return the sequence and state
                                   return_sequences=True,
                                   return_state=True,
                                   recurrent_initializer='glorot_uniform')

  def call(self, tokens, state=None):
    # 2. The embedding layer looks up the embedding for each token.
    vectors = self.embedding(tokens)

    # 3. The GRU processes the embedding sequence.
    #    output shape: (batch, s, enc_units)
    #    state shape: (batch, enc_units)
    output, state = self.gru(vectors, initial_state=state)

    # 4. Returns the new sequence and its state.
    return output, state
```

### Decoder
Decoderの全体部分は後回しにして、Attentionから説明します。
#### Attention
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/134273/0c530bfc-c14e-a2ad-36e1-b957f17714a1.png)
このレイヤが今回の本題です。
少し複雑なので関連部分をクローズアップして計算式とShapeを追記しました。Shapeの(E)はEncoder、(D)はDecoderを示しています(Token(D): DecoderのToken数)。添字のsはEncoderのindex、tはDecoderのindexです。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/134273/ff12529e-3b6f-1980-4c86-0bced9d18716.png)

|入出力|Shape|
|:-:|:-:|
|Value Mask: TextVectorization出力->Attention入力|BS * Token数(E)|
|Value: RNN出力(Output)->Attention入力|BS * Token数(E) * RNN Unit数(E)|
|Query: RNN出力(Output)->Attention入力|BS * Token数(D) * RNN Unit数(D)|
|Attention weight: Attention出力|BS * Token数(E) * Token数(D)|
|Context Vector: Attention出力|BS * Token数(E) * RNN Unit数(E)|

##### Attention weight
score部分の計算に[Bahdanau's additive attention](https://arxiv.org/pdf/1409.0473.pdf)を使っています。今回は使いませんが、もう一つ有名なLuong's multiplicative style というのもあるようです。
Attention weightは[チュートリアルのトップ](https://www.tensorflow.org/text/tutorials/nmt_with_attention)にあるように、EncoderのTokenとDecoderのTokenのマトリックスです(バッチサイズは除く)。どのDecoderのTokenからAttentionされているのがどのEncoderのTokenの割合です(合計が1の確率変数)。Attention weightのshapeは「BS * Token数(E) * Token数(D)」になります。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/134273/91b466cd-4e8f-e50f-5b08-920a4d09d101.png)

Attention weightの計算式です。tanhは活性化関数です。score部の$v_{\alpha}^T$は学習パラメータなのが注意点です。

```math
\alpha_{ts} = \frac{\exp(score(h_t, \bar{h}_s))}{\sum_{s'=1}^{S} \exp(score(h_t, \bar{h}_{s'}))}\\
score(h_t, \bar{h}_s)=v_{\alpha}^T \tanh (W_1 h_t + W_2 \bar{h}_s)
```

Attention weightは、$score(h_t, \bar{h}_s)$のsoftmax関数になっています。参考に下記がsoftmax関数。

```math
{y_i = \frac{\exp(x_i)}{\sum_{j=1}^N \exp(x_j)}}
```

内部的にはvalueもkeyもmaskを`tf.keras.layers.AdditiveAttention`に渡しています。maskをどう使っているか細かくは調べていません(単純に計算から除外していると推測しています)。

```python:BahdanauAttention
class BahdanauAttention(tf.keras.layers.Layer):
  def __init__(self, units):
    super().__init__()
    # For Eqn. (4), the  Bahdanau attention
    self.W1 = tf.keras.layers.Dense(units, use_bias=False)
    self.W2 = tf.keras.layers.Dense(units, use_bias=False)

    self.attention = tf.keras.layers.AdditiveAttention()

  def call(self, query, value, mask):
    # From Eqn. (4), `W1@ht`.
    w1_query = self.W1(query)

    # From Eqn. (4), `W2@hs`.
    # from memory to key
    w2_key = self.W2(value)

    # 3つ目の次元(RNN の unit)をなくしてすべてTrueのBool型
    query_mask = tf.ones(tf.shape(query)[:-1], dtype=bool)
    value_mask = mask

    context_vector, attention_weights = self.attention(
        inputs = [w1_query, value, w2_key],
        mask=[query_mask, value_mask],
        return_attention_scores = True,
    )

    return context_vector, attention_weights
```

maskはこんな中身。boolで、Tokenがpaddingされている部分はFalse。

```python
tf.Tensor(
[[ True  True  True ... False False False]
 ...
 [ True  True  True ... False False False]], shape=(128, 32), dtype=bool)
```

##### Context vector
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/134273/ff12529e-3b6f-1980-4c86-0bced9d18716.png)
Context vectorはAttention weightとDecoderのRNN Outputの内積の総和。

|入出力|Shape|
|:-:|:-:|
|Query: RNN出力(Output)->Context Vecotr入力|BS * Token数(D) * RNN Unit数(D)|
|Attention weight: Context Vector入力|BS * Token数(E) * Token数(D)|
|Context Vector: Attention出力|BS * Token数(E) * RNN Unit数(E)|

計算式。

```math
c_t = \sum_s \alpha_{ts} \bar{h}_s
```

コードは、「Attention weight」部分と共通。

##### Decoder 全体
最後にDecoder全体です。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/134273/d9a73196-8dca-ae03-b374-83b748aaeed7.png)

|入出力|Shape|
|:-:|:-:|
|Embedding入力|BS * 最大Token数|
|Embedding->RNN|BS * 最大Token数 * Embedding Unit数|
|RNN出力(State)|BS * RNN Unit数|
|RNN出力(Output)->Attention Vecotr|BS * 最大Token数 * RNN Unit数|
|Context Vector->Attention Vector|BS * 最大Token数 * RNN Unit数|
|Attention Vector->Fully Connected|BS * 最大Token数 * RNN Unit数|
|Fully Connected->Output|BS * 最大Token数 * 最大Vocabrary数|

```python:Decoder Definition
class Decoder(tf.keras.layers.Layer):
  def __init__(self, output_vocab_size, embedding_dim, dec_units):
    super(Decoder, self).__init__()
    self.dec_units = dec_units
    self.output_vocab_size = output_vocab_size
    self.embedding_dim = embedding_dim

    # For Step 1. The embedding layer convets token IDs to vectors
    self.embedding = tf.keras.layers.Embedding(self.output_vocab_size,
                                               embedding_dim)

    # For Step 2. The RNN keeps track of what's been generated so far.
    self.gru = tf.keras.layers.GRU(self.dec_units,
                                   return_sequences=True,
                                   return_state=True,
                                   recurrent_initializer='glorot_uniform')

    # For step 3. The RNN output will be the query for the attention layer.
    self.attention = BahdanauAttention(self.dec_units)

    # For step 4. Eqn. (3): converting `ct` to `at`
    self.Wc = tf.keras.layers.Dense(dec_units, activation=tf.math.tanh,
                                    use_bias=False)

    # For step 5. This fully connected layer produces the logits for each
    # output token.
    self.fc = tf.keras.layers.Dense(self.output_vocab_size)
```

```python:Decoder call
def call(self,
         inputs: DecoderInput,
         state=None) -> Tuple[DecoderOutput, tf.Tensor]:

  # Step 1. Lookup the embeddings
  vectors = self.embedding(inputs.new_tokens)

  # Step 2. Process one step with the RNN
  rnn_output, state = self.gru(vectors, initial_state=state)

  # Step 3. Use the RNN output as the query for the attention over the
  # encoder output.
  context_vector, attention_weights = self.attention(
      query=rnn_output, value=inputs.enc_output, mask=inputs.mask)

  # Step 4. Eqn. (3): Join the context_vector and rnn_output
  #     [ct; ht] shape: (batch t, value_units + query_units)
  context_and_rnn_output = tf.concat([context_vector, rnn_output], axis=-1)

  # Step 4. Eqn. (3): `at = tanh(Wc@[ct; ht])`
  attention_vector = self.Wc(context_and_rnn_output)

  # Step 5. Generate logit predictions:
  logits = self.fc(attention_vector)
 
  return DecoderOutput(logits, attention_weights), state
```

##### Attention Vector
Attention Vecotor部分の補足。Decoderの一部です。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/134273/35f9072c-bf6e-7e54-7719-e8379a0f301e.png)

計算式です。

```math
\alpha_t = f(c_t,h_t) = \tanh (W_c[c_t;h_t])
```

Decoder RNNのoutputと Context Vectorを結合して、密結合層に入力しているのがわかります。

```python
# Step 4. Eqn. (3): Join the context_vector and rnn_output
  #     [ct; ht] shape: (batch t, value_units + query_units)
  context_and_rnn_output = tf.concat([context_vector, rnn_output], axis=-1)

  # Step 4. Eqn. (3): `at = tanh(Wc@[ct; ht])`
  # self.Wcの活性化関数がtanhにしている
  attention_vector = self.Wc(context_and_rnn_output)
```

# 少しだけメモありプログラムGitHubリンク
[TensorFlow公式チュートリアル「Neural machine translation with attention」](https://www.tensorflow.org/text/tutorials/nmt_with_attention)をコピーし、メモしたものを[nmt_with_attention.ipynb](https://github.com/YoheiFukuhara/bert_nlp/blob/main/others/nmt_with_attention.ipynb)としてGitHubに置きました。大したことをメモしてないですが、ないよりマシレベル。