In many NLP tasks, to handle batched processing, we pad the inputs in a given batch to match the length of the longest element in the batch. We also get a corresponding mask. Now, my point of confusion is that I am not sure how to handle the masks correctly. In BERT inspired implementations[0], I find it processed as following (self-attention):
def forward(input_tokens, mask):
"""
input_tokens: batch_size x # input_tokens x hidden_dim
mask: batch_size x # input_tokens
The mask value is 1 for valid input and 0 for paddings
"""
1. Expand the mask to batch_size x #input_tokens x hidden_dim
2. Make the padded tokens 0: input_tokens[~expanded_mask] = 0
3. input_tokens = input_tokens + positional_embeddings
4. .. norm layer
5. For each attention layer:
input_tokens = attention(input_tokens, mask)
6. return input_tokens
And then, inside the attention blocks:
def forward(query, key, value, mask):
1. score = query x key / sqrt(dim)
2. score = score.masked_fill(~mask, -float('Inf'))
3. attention = softmax(score)
4. context = attention * value
...
Now, in the above pseudo-code, I have applied masking to the input tokens before feeding it to the attention layers [also mentioned by 1]. And I am also applying it when calculating the scores. Is this the right way to do it?
My confusion arises when I see other implementations online where they don't mask the input tokens before processing and handle masks only during scoring [2].
Am kinda stuck on it. Any help is deeply appreciated. Thank you.
[0] https://github.com/huggingface/transformers/blob/v4.45.1/src/transformers/models/hubert/modeling_hubert.py#L1227 (Refer HubertAttention and HubertEncoder classes)
[1] https://ai.stackexchange.com/a/41133/87646
[2] https://pi-tau.github.io/posts/transformer/#multi-head-attention-layer
