According to https://pytorch.org/docs/stable/generated/torch.nn.TransformerDecoder.html, the forward function of nn.TransformerDecoder contemplates the following arguments:
- tgt – the sequence to the decoder (required).
- memory – the sequence from the last layer of the encoder (required).
- tgt_mask – the mask for the tgt sequence (optional).
- memory_mask – the mask for the memory sequence (optional).
- tgt_key_padding_mask – the mask for the tgt keys per batch (optional).
- memory_key_padding_mask – the mask for the memory keys per batch (optional).
Unfortunately, Pytorch's official documentation on the function isn't exactly very thorough at this point (April 2021), in terms of the expected dimensions of each tensor and when it does or doesn't make sense to use each of the optional arguments.
For example, in previous conversations it was explained to me that tgt_mask is usually a square matrix used for self attention masking to prevent future tokens from leaking into the prediction of past tokens. Similarly, tgt_key_padding_mask is used for masking padding tokens (which happens when you pad a batch of sequences of different lengths so that they can fit into a single tensor). In light of this, it makes total sense to use tgt_mask in the decoder, but I wouldn't be so sure about tgt_key_padding_mask. What would be the point of masking target padding tokens? Isn't it enough to simply ignore the predictions associated to padding tokens during training (say, you could do something like
nn.CrossEntropyLoss(ignore_index=PADDING_INDEX) and that's it)?
More generally, and considering that the current documentation is not as thorough as one would like it to be, I would like to know what the purpose is of each argument of nn.TransformerDecoder's forward function, when it makes sense to use each of the optional arguments, and if there are nuances in the usage one should keep in mind when switching between training and inference modes.