[Transformer入門]TensorFlowチュートリアル: Transformer

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

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

以下のTensorFlowのチュートリアルを使って学習しました。学習途中にチュートリアルが更新され、この記事と一部違う点があります。

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

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

https://qiita.com/omiita/items/07e69aef6c156d23c538

https://qiita.com/norihitoishida/items/2fead107792b504eaccf

https://zenn.dev/attentionplease/articles/2d4b2b55ba396e

https://www.slideshare.net/DeepLearningJP2016/dlattention-is-all-you-need

# レイヤ解説
## Attention 解説
モデル全体を解説する前に Attention 部分を解説します。Attentionがモデル内で複数箇所使われているためです。下図の一番左がモデル全体で、その中で(Masked)**Multi-Head Attention**が使われ、その中で**Scaled Dot-Product Attention**が使われています。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/134273/cbfac275-0f92-8896-2062-eb8ae045631a.png)

### Scaled Dot-Product Attention
最小単位のScaled Dot-Product Attentionから解説します。
[記事「[Attention入門]seq2seqとAttentionの解説(TensorFlow)」](https://qiita.com/FukuharaYohei/items/350ac9ff75671e34cbae)で使ったのは**Bahdanau's additive attention**で、今回使うのは**dot-product attention**(内積Attention)。
下図モデルで、QKV(Query, Key, Value)のインプットを受けるのはadditive attentionと変わりません。**QKYのそれぞれが何を示すかはAttentionがどこで使われるかに依存**します。
Scaled Dot-Product Attentionの処理図です(元の図にAttention Weightを加えました)。QueryとKeyからAttetion Weightを作成し、Valueと乗算して焦点を当てたいTokenを出力をします。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/134273/08ba6e51-ca94-179c-f3d0-c9453439b1c0.png)
dot-product attention の数式です。 $\sqrt{d_k}$ で除算しているので**Scaled** Dot-Product Attentionと呼びます。この数式に"Mask(opt.)"処理は含んでいません。

```math
\mathrm{Attention}(Q, K, V)=\mathrm{softmax}_k(\frac{QK^T}{\sqrt{d_k}})V \\
```

- Query: 次元$d_k$(総要素数)を持つ
- Key: 次元$d_k$(総要素数)を持つ
- Value: 次元$d_v$(総要素数)を持つ

> 内積の注意は、深さの平方根の係数でスケーリングされます。これは、深さの値が大きい場合、ドット積の大きさが大きくなり、勾配が小さいソフトマックス関数を押して、非常にハードなソフトマックスになるためです。
たとえば、 QとKの平均が0で分散が1であるとします。これらの行列の乗算は、平均が0で分散がdkになります。したがって、 dkの平方根がスケーリングに使用されるため、 dkの値に関係なく一貫した分散が得られます。分散が低すぎる場合、出力がフラットすぎて効果的に最適化できない可能性があります。分散が高すぎると、初期化時にソフトマックスが飽和し、学習が困難になる可能性があります。

上記は[チュートリアル「Scaled dot product attention」](https://www.tensorflow.org/text/tutorials/transformer#scaled_dot_product_attention)内の日本語説明文です。分散が大きいと学習がすすまないのでScaleをしていることの説明です。

Pythonでガウス分布の配列を作成して分散について確認してみました。

```python:分散確認
import math
import numpy as np

def check_dot(d1, d2):
    size = d1*d2
    a = np.random.normal(size=size).reshape(d1, d2)
    b = np.random.normal(size=size).reshape(d1, d2)
    dot = np.dot(a, b.T)
    if d1*d2 < 10:
        print(f'a: \n{a}')
        print(f'b: \n{b}')
        print(f'dot product: \n{dot}')
    print(f'average: {dot.mean()}')
    print(f'variance: {dot.var()}')
    print(f'scaled variance: {dot.var()/math.sqrt(size)}')
```

総次元数で正規化しないと分散が大きくなっていくことがわかります。

```python: 分散確認結果
> check_dot(2, 3)
a: 
[[-1.00909328 -0.40230197  0.32471939]
 [ 1.41880302 -1.34959901 -1.332655  ]]
b: 
[[ 0.51595355 -0.65731419  0.10551541]
 [-1.18201672 -0.52490031  0.52338299]]
dot product: 
[[-0.22194357  1.57388616]
 [ 1.47853139 -1.6661329 ]]
average: 0.29108526871573576
variance: 1.7873769938079478
scaled variance: 0.7296936021365334

> check_dot(10, 10)
average: 0.08405168619719078
variance: 5.157999880963599
scaled variance: 0.51579998809636

> check_dot(100, 100)
average: -0.11233147437872655
variance: 99.42723752685359
scaled variance: 0.9942723752685358
```

> The mask is multiplied with -1e9 (close to negative infinity). This is done because the mask is summed with the scaled matrix multiplication of Q and K and is applied immediately before a softmax. The goal is to zero out these cells, and large negative inputs to softmax are near zero in the output.

Maskは-1e9(-10億)でSoftmax前に乗算しています。Softmaxは$\frac{\exp(x_i)}{\sum_{k=1}^n \exp(x_k)}$なので-1e9を乗算することで結果はほぼ0です。Softmaxとその微分は、[こちらのページ](https://atmarkit.itmedia.co.jp/ait/articles/2004/08/news016.html)がわかりやすいです。

Attention部分のコードです。コードが非常にわかりやすいです。

```python
def scaled_dot_product_attention(q, k, v, mask):
  """Calculate the attention weights.
  q, k, v must have matching leading dimensions.
  k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.
  The mask has different shapes depending on its type(padding or look ahead)
  but it must be broadcastable for addition.

  Args:
    q: query shape == (..., seq_len_q, depth)
    k: key shape == (..., seq_len_k, depth)
    v: value shape == (..., seq_len_v, depth_v)
    mask: Float tensor with shape broadcastable
          to (..., seq_len_q, seq_len_k). Defaults to None.

  Returns:
    output, attention_weights
  """

  matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)

  # scale matmul_qk
  dk = tf.cast(tf.shape(k)[-1], tf.float32)
  scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

  # add the mask to the scaled tensor.
  if mask is not None:
    scaled_attention_logits += (mask * -1e9)

  # softmax is normalized on the last axis (seq_len_k) so that the scores
  # add up to 1.
  attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)

  output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)

  return output, attention_weights
```

### Multi-head attention
Transformer内でMasked Multi-Head Attention と「Masked」がついていることもあるAttentionです。「Masked」の有無は実質的にはScaled Dot-Product Attention内での処理の差異です。論文内の図に入出力のShapeなど少し補足しました。「BS」はバッチサイズの略。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/134273/087c52b3-9134-1656-c1bd-6eebbe975403.png)

|処理|入力Shape|出力Shape|
|:-:|:--|:--|
|1.1. 線形(Dense)|VKQ各々共通: BS * Token数|VKQ各々共通: BS * Token数 * 次元数|
|1.2. 線形(reshape & transpose)|VKQ各々共通: BS * Token数 * 次元数|VKQ各々共通: BS * ヘッド数 * Token数 * 深さ
|2. Dot-Product Attention|BS * ヘッド数 * Token数 * 深さ|Scaled Attention: BS * ヘッド数 * Token数(Q) * 深さ <br>Attention Weight: BS * ヘッド数 * Token数(Q) * Token数(K)|
|3. 結合|Scaled Attention: BS * ヘッド数* Token数(Q) * 深さ|BS *  Token数(Q) * 深さ|
|4. 線形|BS *  Token数(Q) * 深さ|BS *  Token数(Q) * 深さ|

コード部分は非常にわかりやすいです。Keras v2.9で[MultiHeadAttention](https://www.tensorflow.org/api_docs/python/tf/keras/layers/MultiHeadAttention)が出来ました。実装する場合には使いましょう。理解に役立つので、[MultiHeadAttention](https://www.tensorflow.org/api_docs/python/tf/keras/layers/MultiHeadAttention)使わないコードも残しておきます。

```python: Multi-head attention
class MultiHeadAttention(tf.keras.layers.Layer):
  def __init__(self,*, d_model, num_heads):
    super(MultiHeadAttention, self).__init__()
    self.num_heads = num_heads
    self.d_model = d_model

    # 割り切れる数値になっていることを確認: %は剰余演算
    assert d_model % self.num_heads == 0

    # モデルの次元 = ヘッド数 * 深さ: //は整数除算(切捨除算) 
    self.depth = d_model // self.num_heads

    self.wq = tf.keras.layers.Dense(d_model)
    self.wk = tf.keras.layers.Dense(d_model)
    self.wv = tf.keras.layers.Dense(d_model)

    self.dense = tf.keras.layers.Dense(d_model)

  def split_heads(self, x, batch_size):
    """Split the last dimension into (num_heads, depth).
    Transpose the result such that the shape is (batch_size, num_heads, seq_len, depth)
    """
    x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth)) # 次で次元順番変更していることに注意
    return tf.transpose(x, perm=[0, 2, 1, 3]) 

  def call(self, v, k, q, mask):
    batch_size = tf.shape(q)[0]

    q = self.wq(q)  # (batch_size, seq_len, d_model)
    k = self.wk(k)  # (batch_size, seq_len, d_model)
    v = self.wv(v)  # (batch_size, seq_len, d_model)

    q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)
    k = self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)
    v = self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)

    # scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth)
    # attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)
    scaled_attention, attention_weights = scaled_dot_product_attention(
        q, k, v, mask)

    scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])  # (batch_size, seq_len_q, num_heads, depth)

    concat_attention = tf.reshape(scaled_attention,
                                  (batch_size, -1, self.d_model))  # (batch_size, seq_len_q, d_model)

    output = self.dense(concat_attention)  # (batch_size, seq_len_q, d_model)

    return output, attention_weights
```

## Point wise feed forward network
Attention とともにモデル構成要素の1つである**Point wise feed forward network**についても解説します。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/134273/8a1015e1-8f84-fc4c-733c-4ab5750b0b01.png)

コードは`Dense`を2つ重ねているだけなので非常にシンプルです。

```python: Point wise feed forward network
def point_wise_feed_forward_network(d_model, dff):
# ここでdffはdimension of feed forward
  return tf.keras.Sequential([
      # 一度、dffで次元を増やす(直後に元のd_modelに戻す)
      tf.keras.layers.Dense(dff, activation='relu'),  # (batch_size, seq_len, dff)
      tf.keras.layers.Dense(d_model)  # (batch_size, seq_len, d_model)
  ])
```

## Positional Encoding
レイヤ説明の最後としてPositional Encoding(位置エンコーディング)です。
EncoderおよびDecoderに使われています。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/134273/4a569579-98f5-3706-4cf9-51938ac37fc2.png)

Positional Encodingの元ネタは以下の論文です。

https://arxiv.org/abs/1705.03122

Positional Encodingの計算式です。入力内容(文字)に関係ない固定値を作成し、加算します(加算部分はコードに記述なし)。

```math
\Large{PE_{(pos, 2i)} = \sin(pos / 10000^{2i / d_{model}})} \\
\Large{PE_{(pos, 2i+1)} = \cos(pos / 10000^{2i / d_{model}})}
```

コードとともに見ていきます。
入力のlengthはTokenの最大長、depthは直前のEmbeddingから受け取る次元数です。`depth`を2で割っているのは、`sin`と`cos`の計算で2回同じ値を使うためです。

```python
def positional_encoding(length, depth):
  depth = depth/2

  # Token Length の配列([[0, ,1, ..., length]])
  # np.newaxisはNoneのエイリアスで、次元を増やす場合に使用
  positions = np.arange(length)[:, np.newaxis]     # (seq, 1)
  print(f'pos: {positions}')

  # Dimension の配列([[0], [2/(depth)],...,[(depth-2)/depth]])
  depths = np.arange(depth)[np.newaxis, :]/depth   # (1, depth/2)
  print(f'dep1: {np.arange(depth)[np.newaxis, :]}')
  print(f'dep2: {depths}')
  
  angle_rates = 1 / (10000**depths)         # (1, depth/2)
  print(f'angle_rates: {angle_rates}')

  angle_rads = positions * angle_rates      # (pos, depth/2)
  print(f'angle_rads: {angle_rads}')

  pos_encoding = np.concatenate(
      [np.sin(angle_rads), np.cos(angle_rads)],
      axis=-1) 

  return tf.cast(pos_encoding, dtype=tf.float32)
```

細かく解説します。
試しにlength=8, dpeth=10で呼び出します。

```python
pos_encoding = positional_encoding(length=8, depth=10)
```

各計算式と値の推移です。

```python
>> np.arange(length)[:, np.newaxis]
[[0] [1] [2] [3] [4] [5] [6] [7]]

>> np.arange(depth)[np.newaxis, :]
[[0. 1. 2. 3. 4.]]

>> np.arange(depth)[np.newaxis, :]/depth
[[0.  0.2 0.4 0.6 0.8]]

>> 1 / (10000**depths)
[[1.00000000e+00 1.58489319e-01 2.51188643e-02 3.98107171e-03
  6.30957344e-04]]
```

表形式。`depth`のindexが多くなるにつれて1から0に近づいていくのがわかります。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/134273/e7f213c0-da54-600d-d171-035f70b51e09.png)

行列計算をすることで、`position/2`×`depth`のshapeになります。

```python
positions * angle_rates
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/134273/f92f043f-6ff8-862a-c4d8-18498460a609.png)

最後に三角関数で周期化します。`depth`(行)のsinの後ろ(3・4行目)はほぼ0付近で、cosの後ろ(8・9・10行目)はほぼ1付近に値が張り付いています。

```python
 pos_encoding = np.concatenate(
      [np.sin(angle_rads), np.cos(angle_rads)],
      axis=-1) 
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/134273/4bdd0221-f60b-1482-7b72-05a9cf57b1b5.png)

[三角関数](https://ja.wikipedia.org/wiki/%E4%B8%89%E8%A7%92%E9%96%A2%E6%95%B0)の復習としてグラフを表示しておきます。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/134273/9390fc4d-cd8b-00ee-f078-9ae434d3f79f.png)

元の`depth`512と`position`2048をグラフ表示するとこんな。`Depth`が0から100や256から356くらいは-1から1までまばらになっているのがわかります。これは周期関数である`sin`や`cos`を使っているからですね。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/134273/de9d6a59-1c41-8c36-2700-2cb0e35dd0d1.png)

Embeddingと繋げて`PositionalEmbedding`というレイヤを作ります。Embeddingとの間にある`tf.math.sqrt(tf.cast(self.d_model, tf.float32))`は論文内でも以下の説明しかありません。Embedding結果に対してPositional Encoding結果を加算をしているのがわかります。

> In the embedding layers, we multiply those weights by $\sqrt{d_{model}}$.

```python
class PositionalEmbedding(tf.keras.layers.Layer):
  def __init__(self, vocab_size, d_model):
    super().__init__()
    self.d_model = d_model
    self.embedding = tf.keras.layers.Embedding(vocab_size, d_model, mask_zero=True) 
    self.pos_encoding = positional_encoding(length=2048, depth=d_model)

  #マスクを作成
  def compute_mask(self, *args, **kwargs):
    return self.embedding.compute_mask(*args, **kwargs)

  def call(self, x):
    length = tf.shape(x)[1]
    x = self.embedding(x)
    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
    x = x + self.pos_encoding[tf.newaxis, :length, :]
    return x
```

# モデル解説
モデル全体の流れは、以下の記事が非常にわかりやすいです。具体例を元に解説してくれています。
逆に私はもう解説することあまりありません。

https://qiita.com/FuwaraMiyasaki/items/239f3528053889847825#attention

## Encoder
Masked Mulit-Head Attentionでないのに、なぜ`attention_mask`があるのか一瞬混乱しました。ここのMaskは`<PAD>`部分なので問題ないです。Masked Mulit-Head Attentionは、カンニングを防ぐためにDecoderで使います。
[`MultiHeadAttention` Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/MultiHeadAttention)がv2.9で加わって便利になりました。

```python
class EncoderLayer(tf.keras.layers.Layer):
  def __init__(self,*,
               d_model, # Input/output dimensionality.
               num_attention_heads,
               dff, # Inner-layer dimensionality.
               dropout_rate=0.1
               ):
    super().__init__()


    # Multi-head self-attention.
    self.mha = tf.keras.layers.MultiHeadAttention(
        num_heads=num_attention_heads,
        key_dim=d_model, # Size of each attention head for query Q and key K.
        dropout=dropout_rate,
        )
    # Point-wise feed-forward network.
    self.ffn = point_wise_feed_forward_network(d_model, dff)

    # Layer normalization.
    self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)

    # Dropout for the point-wise feed-forward network.
    self.dropout1 = tf.keras.layers.Dropout(dropout_rate)

  def call(self, x, training, mask):

    # A boolean mask.
    if mask is not None:
      mask1 = mask[:, :, None]
      mask2 = mask[:, None, :]
      attention_mask = mask1 & mask2
    else:
      attention_mask = None

    # Multi-head self-attention output (`tf.keras.layers.MultiHeadAttention `).
    attn_output = self.mha(
        query=x,  # Query Q tensor.
        value=x,  # Value V tensor.
        key=x,  # Key K tensor.
        attention_mask=attention_mask, # A boolean mask that prevents attention to certain positions.
        training=training, # A boolean indicating whether the layer should behave in training mode.
        )

    # Multi-head self-attention output after layer normalization and a residual/skip connection.
    out1 = self.layernorm1(x + attn_output)  # Shape `(batch_size, input_seq_len, d_model)`

    # Point-wise feed-forward network output.
    ffn_output = self.ffn(out1)  # Shape `(batch_size, input_seq_len, d_model)`
    ffn_output = self.dropout1(ffn_output, training=training)
    # Point-wise feed-forward network output after layer normalization and a residual skip connection.
    out2 = self.layernorm2(out1 + ffn_output)  # Shape `(batch_size, input_seq_len, d_model)`.

    return out2
```

```python
class Encoder(tf.keras.layers.Layer):
  def __init__(self,
               *,
               num_layers,
               d_model, # Input/output dimensionality.
               num_attention_heads,
               dff, # Inner-layer dimensionality.
               input_vocab_size, # Input (Portuguese) vocabulary size.
               dropout_rate=0.1
               ):
    super().__init__()

    self.d_model = d_model
    self.num_layers = num_layers

    # Embeddings + Positional encoding
    self.pos_embedding = PositionalEmbedding(input_vocab_size, d_model)

    # Encoder layers.
    self.enc_layers = [
        EncoderLayer(
          d_model=d_model,
          num_attention_heads=num_attention_heads,
          dff=dff,
          dropout_rate=dropout_rate)
        for _ in range(num_layers)]
    # Dropout.
    self.dropout = tf.keras.layers.Dropout(dropout_rate)

  # Masking.
  def compute_mask(self, x, previous_mask=None):
    # 2つ目のパラメータ previous_mask は無意味
    # https://github.com/keras-team/keras/blob/v2.10.0/keras/layers/core/embedding.py#L173
    return self.pos_embedding.compute_mask(x, previous_mask)

  def call(self, x, training):

    seq_len = tf.shape(x)[1]

    # Sum up embeddings and positional encoding.
    mask = self.compute_mask(x)
    x = self.pos_embedding(x)  # Shape `(batch_size, input_seq_len, d_model)`.
    # Add dropout.
    x = self.dropout(x, training=training)

    # N encoder layers.
    for i in range(self.num_layers):
      x = self.enc_layers[i](x, training, mask)

    return x  # Shape `(batch_size, input_seq_len, d_model)`.
```

## Decoder
ここの1つ目のAttentionがMasked Mulit-Head Attentionで、`use_causal_mask=True`としています。これによってカンニングを防ぎます。

```python
class DecoderLayer(tf.keras.layers.Layer):
  def __init__(self,
               *,
               d_model, # Input/output dimensionality.
               num_attention_heads,
               dff, # Inner-layer dimensionality.
               dropout_rate=0.1
               ):
    super().__init__()

    # Masked multi-head self-attention.
    self.mha_masked = tf.keras.layers.MultiHeadAttention(
        num_heads=num_attention_heads,
        key_dim=d_model, # Size of each attention head for query Q and key K.
        dropout=dropout_rate
    )
    # Multi-head cross-attention.
    self.mha_cross = tf.keras.layers.MultiHeadAttention(
        num_heads=num_attention_heads,
        key_dim=d_model, # Size of each attention head for query Q and key K.
        dropout=dropout_rate
    )

    # Point-wise feed-forward network.
    self.ffn = point_wise_feed_forward_network(d_model, dff)

    # Layer normalization.
    self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    self.layernorm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)

    # Dropout for the point-wise feed-forward network.
    self.dropout1 = tf.keras.layers.Dropout(dropout_rate)

  def call(self, x, mask, enc_output, enc_mask, training):
    # The encoder output shape is `(batch_size, input_seq_len, d_model)`.

    # A boolean mask.
    self_attention_mask = None
    if mask is not None:
      mask1 = mask[:, :, None]
      mask2 = mask[:, None, :]
      self_attention_mask = mask1 & mask2

    # Masked multi-head self-attention output (`tf.keras.layers.MultiHeadAttention`).
    attn_masked, attn_weights_masked = self.mha_masked(
        query=x,
        value=x,
        key=x,
        attention_mask=self_attention_mask,  # A boolean mask that prevents attention to certain positions.
        use_causal_mask=True,  # A boolean to indicate whether to apply a causal mask to prevent tokens from attending to future tokens.
        return_attention_scores=True,  # Shape `(batch_size, target_seq_len, d_model)`.
        training=training  # A boolean indicating whether the layer should behave in training mode.
        )

    # Masked multi-head self-attention output after layer normalization and a residual/skip connection.
    out1 = self.layernorm1(attn_masked + x)

    # A boolean mask.
    attention_mask = None
    if mask is not None and enc_mask is not None:
      mask1 = mask[:, :, None]
      mask2 = enc_mask[:, None, :]
      attention_mask = mask1 & mask2

    # Multi-head cross-attention output (`tf.keras.layers.MultiHeadAttention `).
    attn_cross, attn_weights_cross = self.mha_cross(
        query=out1,
        value=enc_output,
        key=enc_output,
        attention_mask=attention_mask,  # A boolean mask that prevents attention to certain positions.
        return_attention_scores=True,  # Shape `(batch_size, target_seq_len, d_model)`.
        training=training  # A boolean indicating whether the layer should behave in training mode.
    )

    # Multi-head cross-attention output after layer normalization and a residual/skip connection.
    out2 = self.layernorm2(attn_cross + out1)  # (batch_size, target_seq_len, d_model)

    # Point-wise feed-forward network output.
    ffn_output = self.ffn(out2)  # Shape `(batch_size, target_seq_len, d_model)`.
    ffn_output = self.dropout1(ffn_output, training=training)
    out3 = self.layernorm3(ffn_output + out2)  # Shape `(batch_size, target_seq_len, d_model)`.

    return out3, attn_weights_masked, attn_weights_cross
```

```python
class Decoder(tf.keras.layers.Layer):
  def __init__(self,
               *,
               num_layers,
               d_model, # Input/output dimensionality.
               num_attention_heads,
               dff, # Inner-layer dimensionality.
               target_vocab_size,
               dropout_rate=0.1
               ):
    super().__init__()

    self.d_model = d_model
    self.num_layers = num_layers

    self.pos_embedding = PositionalEmbedding(target_vocab_size, d_model)

    self.dec_layers = [
        DecoderLayer(
          d_model=d_model,
          num_attention_heads=num_attention_heads,
          dff=dff,
          dropout_rate=dropout_rate)
        for _ in range(num_layers)
    ]
    self.dropout = tf.keras.layers.Dropout(dropout_rate)

  def call(self, x, enc_output, enc_mask, training):
    attention_weights = {}

    mask = self.pos_embedding.compute_mask(x)
    x = self.pos_embedding(x)  # Shape: `(batch_size, target_seq_len, d_model)`.

    x = self.dropout(x, training=training)

    for i in range(self.num_layers):
      x, block1, block2  = self.dec_layers[i](x, mask, enc_output, enc_mask, training)

      attention_weights[f'decoder_layer{i+1}_block1'] = block1
      attention_weights[f'decoder_layer{i+1}_block2'] = block2

    # The shape of x is `(batch_size, target_seq_len, d_model)`.
    return x, attention_weights
```

## Transformer
最後にTransformer全体です。

```python
class Transformer(tf.keras.Model):
  def __init__(self,
               *,
               num_layers, # Number of decoder layers.
               d_model, # Input/output dimensionality.
               num_attention_heads,
               dff, # Inner-layer dimensionality.
               input_vocab_size, # Input (Portuguese) vocabulary size.
               target_vocab_size, # Target (English) vocabulary size.
               dropout_rate=0.1
               ):
    super().__init__()
    # The encoder.
    self.encoder = Encoder(
      num_layers=num_layers,
      d_model=d_model,
      num_attention_heads=num_attention_heads,
      dff=dff,
      input_vocab_size=input_vocab_size,
      dropout_rate=dropout_rate
      )

    # The decoder.
    self.decoder = Decoder(
      num_layers=num_layers,
      d_model=d_model,
      num_attention_heads=num_attention_heads,
      dff=dff,
      target_vocab_size=target_vocab_size,
      dropout_rate=dropout_rate
      )

    # The final linear layer.
    self.final_layer = tf.keras.layers.Dense(target_vocab_size)

  def call(self, inputs, training):
    # Keras models prefer if you pass all your inputs in the first argument.
    # Portuguese is used as the input (`inp`) language.
    # English is the target (`tar`) language.
    inp, tar = inputs

    # The encoder output.
    enc_output = self.encoder(inp, training)  # `(batch_size, inp_seq_len, d_model)`
    enc_mask = self.encoder.compute_mask(inp)

    # The decoder output.
    dec_output, attention_weights = self.decoder(
        tar, enc_output, enc_mask, training)  # `(batch_size, tar_seq_len, d_model)`

    # The final linear layer output.
    final_output = self.final_layer(dec_output)  # Shape `(batch_size, tar_seq_len, target_vocab_size)`.

    # Return the final output and the attention weights.
    return final_output, attention_weights
```