I am trying to implement an autoregressive transformer model similar to the paper attention is all you need. From what I have understood, in order to replicate the architecture fully, I need to give the transformer decoder 3 masks.

  1. Target subsequent mask: this is for causality.

  2. Target padding indexes: just to look at non-padded indices.

  3. Encoder padding indices: just to look at non-padded inputs from the encoder.

The snippet is here:

y = self.decoder(y, x,

With masks being generated like this:

def generate_no_peek_mask(self, sz):
    mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
    mask = mask.float().masked_fill(mask == 0, float("-inf")).masked_fill(mask == 1, float(0.0))
    mask = mask.to(self.device)
    return mask

def generate_padding_mask(self, seq, pad_idx):
    return (seq != pad_idx).to(self.device)

The problem is that using these masks leads to issues with the Softmax function because of NaN values. Without these masks, the model does not generate any NaN value. I have tried toying with various input lengths and seeing what happens when I make sure my inputs are moderately big, but it still does not work. The only thing that works is not giving the decoder the masks.


1 Answer 1


I found out what the problem was. It was not from subsequent mask. It was caused by bad key_padding_mask. PyTorch expects the key_padding_mask to be true wherever there is a padding token and false wherever there is none. I was generating the padding mask in exactly the opposite way. So the correct mask will be:

def generate_padding_mask(self, seq, pad_idx):
    return (seq == pad_idx).to(self.device)

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

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