I'd like to know, from an implementation point of view, when training a Causal Transformer such as GPT-2, if making prediction on whole sequence at once and computing the loss on the whole sequence is something standard ?
When going across examples that vulgarize what happen during training, we have examples like this:
Which suggests that we mask at a certain token in the sequence and make a prediction, then compute the loss for this single token, so the loss would take data of shape
However if I'm correct, since we're talking about causal models, we could predict all tokens at once because the model can only attend to what's on the left of the sequence (and can't attend on the right, so it can't "cheat"), apply the loss on data that would have the shape
(batch_size, sequence_length, num_classes) where
sequence_length is computed in a single forward pass. And so speedup the training.
Am I correct ? If so, do you know famous repos that do this ? And if not, why would it be wrong ?