Loss on whole sequences in Causal Language Model

Question

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 (batch_size, num_classes).

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 ?

Thanks.

Answer 1

The figure and the blog post are simply incorrect. Doing a reverse image search, I see that the image you posted comes from a blog post on Towards Data Science. That image is so wrong. Just think that in a causal language model, the prediction for the word next would be in the time step that receives as input the word very! Also, causal language models do not use any <mask> tokens.

You are correct in your understanding: at training time, the loss of all time steps is computed at once. This can be done because:

• The models attend only to the previous tokens. RNNs do this by construction. Transformers do this because of the mask imposed on the decoder.
• We are using the true previous tokens as input to the prediction instead of using the model's own predictions as input. This is referred to as "teacher forcing".

ALL causal language modeling implementations compute the training loss in a single pass.