How to understand the self-attention mask implementation in google transformer tutorial

Viewed 11

I am reading google's transformer tutorial, and the part why the attention_mask for multi-head attention can be built via mask1 & mask2 was unclear to me. Any help would be great!

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

    # A boolean mask.
    if mask is not None:
      mask1 = mask[:, :, None]
      mask2 = mask[:, None, :]
      attention_mask = mask1 & mask2          # <= here 
    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.
        )

toy example breakdown

input = tf.constant([
    [[1, 0, 3, 0], [1, 2, 0, 0]]
])

mask = tf.keras.layers.Embedding(2,2, mask_zero=True).compute_mask(input)
print(mask)
mask1 = mask[:, :, None]   # same as tf.expand_dims(mask, axis = 2)
print(mask1)
mask2 = mask[:, None, :]
print(mask2)

print(mask1 & mask2)

>

tf.Tensor(
[[[ True False  True False]
  [ True  True False False]]], shape=(1, 2, 4), dtype=bool)

tf.Tensor(
[[[[ True False  True False]]

  [[ True  True False False]]]], shape=(1, 2, 1, 4), dtype=bool)

tf.Tensor(
[[[[ True False  True False]
   [ True  True False False]]]], shape=(1, 1, 2, 4), dtype=bool)

<tf.Tensor: shape=(1, 2, 2, 4), dtype=bool, numpy=               # <= why built mask like this?
array([[[[ True, False,  True, False],
         [ True, False, False, False]],

        [[ True, False, False, False],
         [ True,  True, False, False]]]])>
0 Answers
Related