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 "perfect" result?