When do we apply a mask onto our padded values during attention mechanisms

Question

When we are applying a mask onto the padded values in an input sequence, it is typically done through setting the padded values as negative infinity. For example, a tensor of values [1,2,3,0,0] should result in a padding mask of pad_mask = [True, True, True, False, False] (or the opposite depending on your flavour). However, if we apply the mask i.e attention_scores = attention_scores.masked_fill_(pad_mask.T == False, float('-inf')) before applying softmax, won't we get the 4th and 5th row of the attention_scores as 'nan' when we softmax attempts to calculate the probability distribution along each row?

Does that mean the step of where to apply the mask is incorrect, and we should apply a zero-ing out of the pad token rows in the attention_score matrix after applying the softmax function? or is there another key concept/step I am missing here

Answer 1

You are masking at the wrong place. Masking happens before the sequence goes into the encoder/decoder layer depending on what kind of architecture you are using. It happens right after you calculate the embeddings using positional encoding + token encodings.

This is because the masked positions are also considered in calculating the attention values, and if you simply make those masked attention values zero, the model loses the attention and can't make accurate predictions.

Answer 2

This is the correct place to use the mask. If you are concerned that applying the softmax afterwards will produce nan, then instead of float(-inf) just use some relatively large number, e.g. -1e9:

attention_scores = attention_scores.masked_fill_(pad_mask.T == False, -1e9)

Why not mask after the softmax ?

After you apply the softmax function to the attention scores you get the attention weights (also called attention probabilities). In general, you want these to sum up to $1.$, and this is exactly what you get with the softmax layer. However, if you mask some attention weights after the softmax, then they will no longer sum to $1.$

Setting the attention scores to a high negative number will effectively make the softmax produce $0.$ at these positions and will re-distribute the masked weight over to the un-masked positions so that the attention weights sum to $1.$

You can read more (and see some cool pictures) here.