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]]]])>