I'm trying to implement transformer model using this tutorial. In the decoder block of the Transformer model, a mask is passed to "pad and mask future tokens in the input received by the decoder". This mask is added to attention weights.

import tensorflow as tf

def create_look_ahead_mask(size):
    mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)
    return mask

Now my question is, how is doing this step (adding mask to the attention weights) equivalent to revealing the words to model one by one? I simply can't grasp the intuition of it's role. Most tutorials won't even mention this step like it's very obvious. Please help me understand. Thanks.


The Transformer model presented in this tutorial is an auto-regressive Transformer. Which means that prediction of next token only depends on it's previous tokens.

So in order to predict next token, you have to make sure that only previous token are attended. (If not, this would be a cheating because model already knows whats next).

So attention mask would be like this
[0, 1, 1, 1, 1]
[0, 0, 1, 1, 1]
[0, 0, 0, 1, 1]
[0, 0, 0, 0, 1]
[0, 0, 0, 0, 0]

For example: If you are translating English to Spanish
Input: How are you ?
Target: < start > Como estas ? < end >
Then decoder will predict something like this
< start > (it will be given to decoder as initial token)
< start > Como
< start > Como estas
< start > Como estas ?
< start > Como estas ? < end >

Now compare this step by step prediction sequences to attention mask given above, It would make sense now to you

| improve this answer | |
  • $\begingroup$ In spanish double question marks are used ¿? $\endgroup$ – pasaba por aqui Oct 24 at 10:23

We give the target input into the transformer decoder while training the model. So it is easy for the model to "peek ahead" and learn what the next word would be. To ensure that this doesn't happen we apply an additive mask after the dot product between Query and Key. In the original paper "Attention is all you need", the triangular matrix had 0's in the lower triangle and -10e9 (You can see negative infinity used in recent examples) in the upper triangle. So when the mask is added to the attention score the attention scores in the upper triangle would be really low. When this matrix is passed through the softmax function, these really low values become close to 0, which essentially means not to attend to the words after timestep t. To put in matrix format,

[8.1, 0.04, 5.2, 4.2]
[0.5, 9.2, 2.33, 0.7]
[0.2, 0.4, 6.11, 1.0]
[3.1, 2.1. 2.19, 8.1]

Let the above matrix A the result of the dot product between query and key. The A[0][0] contains the attention score of the first-word query to the first word of the key, A[0][1] contains the attention score of the first word of the query to the second of the key, and so on. So as you can see the after adding the mask and performing softmax on A, the result would be,

[8.1, 0.0, 0.0, 0.0]
[0.5, 9.2, 0.0, 0.0]
[0.2, 0.4, 6.11, 0.0]
[3.1, 2.1. 2.19, 8.1]

This forces the transformer only to attend to words that are before it. You can check out the Transformer lecture available in CS224n for full detail.

| improve this answer | |

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.