Transformers: How to use the target mask properly?

Question

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 = self.decoder(tgt, memory, self.tgt_mask.type_as(tgt))
    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.

Answer 1

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\|$).