# Transformers: How to use the target mask properly?

I try to apply Transformers to an unusual use case - predict the next user session based on the previous one. A user session is described by a list of events per second, e.g. whether the user watches a particular video, clicks a specific button, etc. Typical sessions are around 20-30 seconds, I pad them to 45 seconds. Here's a visual example of 2 subsequent sessions:

x axis is time in seconds, y axis is the list of events (black line divides the 2 sessions). I extend the vocabulary with 2 additional tokens - start and end of a session (<sos> and <eos>), where <sos> is a one-hot vector at the very beginning and <eos> - similar vector at the end of the session (which makes this long red line).

Now I use these extended vectors of events as embeddings and want to train a Transformer model to predict the next events in the current session based on previous events in this (target) session and all events in the previous (source) session. So pretty much like seq2seq autoregressive models, but in a bit unusual settings.

Here's the problem. When I train a Transformer using the built-in PyTorch components and square subsequent mask for the target, my generated (during training) output is too good to be true:

Although there's some noise, many event vectors in the output are modeled exactly as in the target. After checking train-val-test split is correct, my best guess is that the model cheats by attending to the same day in the target, which the mask should have prevented. The mask is (5x5 version for brevity):

[[0., -inf, -inf, -inf, -inf],
[0., -inf, -inf, -inf, -inf],
[0., 0., -inf, -inf, -inf],
[0., 0., 0., -inf, -inf],
[0., 0., 0., 0., -inf]]

Note that since I use <sos> in both - source and target - mask[i, i] is set to -inf (except for mask[0, 0] for numerical reasons), so the output timestamp i should not attend to the target timestamp i.

The code for the model's forward method:

def forward(self, src, tgt):
memory = self.encoder(src)
out = torch.sigmoid(out)
return out

I also tried to avoid the target mask altogether and set it to all -inf (again, except for the first column for numerical stability), but the result is always the same.

Am I using the mask the wrong way? If the mask looks fine, what other reasons could lead to such a "perfect" result?

After shifting the target to the right as suggested in the accepted answer I get the following result:

Which is much more realistic. One suspicious thing is that out[t] now resembles tgt[t - 1], but it can be explained by the fact that the user state tends to be "sticky", e.g. if a user watches a video at t - 1, most likely he will watch it at t as well.

• When you say "I use these extended vectors of events as embeddings", do you mean that you do not learn the initial $\|dict\|$ x $d_{model}$ weight matrix to embed the inputs, but instead feed the inputs raw into the first layer of the transformer? And do you perform any positional encoding? Also, what is the dimension of your model ($d_{model}$ in some papers), and what is the dimension of your mask? – user3667125 Dec 31 '20 at 18:11
• It would also help if we can read the code of your entire Transformer model. Something like the code under "Define the Model" section here (pytorch.org/tutorials/beginner/transformer_tutorial.html). Would help us clarify some details like what is self.encoder, self.decoder, etc. – user3667125 Dec 31 '20 at 18:23
• @user3667125 Here's the code for the whole model. It's a PyTorch Lightning module, but you can safely ignore all methods except for __init__() and forward(). The loss should also be irrelevant. Regarding "embeddings", yes, you got it right: my sequences are not tokens, but instead vectors of events, so I decided to use them directly instead of usual embeddings, without learnable matrix or positional encoding. $d_{model} = 14$, sequences are padded to 45 seconds and thus mask is 45x45. – ffriend Dec 31 '20 at 18:42
• Can you confirm if a single instance have dimension $45$ x $14$ (and therefore a batch has dimension $45$ x $batchSize$ x $14$)? And can you explain why the mask diagonal entries mask[i, i] = -inf? Normally they would be 0. – user3667125 Dec 31 '20 at 20:24
• Also, what is the reasoning behind type-casting the mask? self.tgt_mask.type_as(tgt) -- have you tried without using the .type_as typecasting? – user3667125 Dec 31 '20 at 20:39

The main issue during training is that you haven't right-shifted the input of the decoder, which is probably why you set the diagonals of mask to -inf (when it should be $$0$$).

Also, just an FYI, although you haven't focused on evaluation/prediction yet, I will explain the evaluation/prediction here as well for completeness, since it works so differently than training, and also since you will need it when generating the graphs for debugging.

## Training

Both tgt and tgt_mask should be changed to simulate auto-regressive properties.

You are feeding in tgt as the input to the decoder, where tgt is the ground truth target sequence to predict. tgt should have dimension length(sequence) x batchSize x $$\|dict\|$$. Additionally, you are feeding in mask where the diagonals are -inf.

Instead, you should do the following:

• Add 2 new special tokens to your vocabulary that signals the start and end of decoding, i.e. <START> and <END>. So your vocabulary will need to be extended to 16 events.
• Every input to the decoder (training, evaluation, prediciton) should have <START> as the first token. So during training, when you want to feed in tgt to the decoder, you should right-shift tgt by adding in the <START> token at the beginning.
• tgt_shifted will now be of dimension length(sequence)+1 x batchSize x $$\| dict \|$$
• tgt_shifted_mask will now be of dimension length(sequence)+1 x length(sequence)+1. Diagonal should be all $$0$$
• the inputs to the decoder should be tgt_shifted, tgt_shifted_mask, and memory
• the output of the decoder will have dimension length(sequence)+1 x batchSize x $$\| dict\|$$, and will not be right-shifted because it will not have <START> as the first word. But it should have <END> as the last word.
• Append <END> to the ground truth tgt, and turn it into one-hot encoding, so tgt should have dimension length(sequence)+1 x batchSize x $$\|dict\|$$. Your loss operator should be some sort of element-wise comparison between tgt and the output of the decoder.

Note that you only run the decoder once per training batch during training, i.e. the decoder simultaneously predicts the logits of all length(sequence) tokens at the same time. On the other hand, during evaluation/prediction, you must run the decoder length(sequence) times, since you can only use the decoder to predict one token at a time.

## Evaluation/Prediction

During evaluation/prediction, the model should generate its own sentence entirely from scratch, one word at a time:

1. out should be initialized as a single token, <START>, and have dimensions $$1$$ x batchSize x $$\|dict\|$$.
2. Do the following in a loop until termination or max output length is reached:
1. Generate a out_mask for the current version of out. out_mask should be a matrix with dimension length(out) x length(out). Diagonal should be all $$0$$
2. Pass out, out_mask, and memory into the decoder layer. The output of the decoder will be a length(out) x batchSize x $$d_{model}$$ tensor (note that in your case, $$d_{model} = \|dict\|$$), which indicates tokenIdx x batchIdx x vocabularyIdx
3. Convert the last time step (e.g. decoderOutput[-1,:,:] if all sequences in your batch have the same length) into probabilities, and take all the events that surpasses some probability threshold and set them to $$1$$. Set all other events to $$0$$. This is the new time step that the decoder layer predicts. So add this time step to out to append it's length by $$1$$.
4. Repeat the previous 3 steps in a loop until you reach the <END> token, which signals the end of decoding, or until the max output length is reached.

During evaluation, you can calculate the validation loss by using the logits of only the final iteration of the decoder (which should have dimension length(sequence)+1 x batchSize x $$\|dict\|$$), and compare those logits with the ground truth with the <END> token added at the end (resulting in dimensions length(sequence)+1 x batchSize x $$\|dict\|$$).

• Sorry, I wasn't clear about this - I get this issue specifically during training, not inference time. Here's how I understand training should go: for an output token at timestamp t we give a model the whole src sequence as well as tgt[0 : t-1]. It's not like generating the whole sentence in English given a sentence in French, but instead like predicting the next word user is going to write given previous sentence and previous words in this sentence user has already written. IIUC, for each out[t] mask should hide tgt[t : ] even though training loop sees the whole tgt. – ffriend Jan 2 at 16:16
• But visualization in the question shows that out[t] is extremely similar to tgt[t]. I can think of 2 reasons for this - overfitting or wrong usage of the mask. I have already checked that the visualization is done on the test set, so the problem must be the mask. – ffriend Jan 2 at 16:26
• Perhaps, the question can be reduced to: is decoder(tgt, memory, subsequent_mask) the correct way to predict all outputs in parallel, restricting each out[t] to see only true values tgt[ : t - 1]? – ffriend Jan 2 at 16:33
• Thanks for the clarification. I have edited my answer which should help with all the things you mentioned. Let me know if there's anything further to clarify. – user3667125 Jan 2 at 22:50
• Glad it helped! One more thing I'd verify is why out[t] resembles tgt[t-1] though -- it does seem weird. Wouldn't be surprised if there's an off-by-one somewhere, since we shift things so many times by 1 token. – user3667125 Jan 4 at 2:07