Are transformer decoder predictions computed in parallel during training?

Question

I've been studying the transformer from the original "Attention is all you need" paper and from various other sources. I have a question about the behaviour of the decoder during training that I cannot find the answer to anywhere.

During inference I understand that the decoder input is its own previously generated token from the prior time step. Tokens are fed into the decoder one-by-one and predictions made one-by-one.

However, during training the target sequence is known and I have read from several sources that the entire sequence is used as the decoder input to allow parallel processing and to improve training efficiency. To keep the decoder autoregressive, a masked attention sub-layer is introduced where a masking matrix is added to the scaled dot product attention mechanism.

So my question is, since during training the decoder input is the entire sequence, is the entire output sequence predicted in parallel (simultaneously), or are tokens predicted one-by-one, as in inference?

To me it makes sense that if the entire target sequence is used as the decoder input, then an entire sequence is output. If it wasn't, the decoder would be using the same input at every timestep whilst being expected to produce different tokens.

Answer 1

In short, masking is just sequence padding in the decoder!

I'll try to slightly rephrase your question first to ensure I properly understood it. You are confused about:

1. The transformer is fed the whole sequence (during training) for efficiency,
2. But it needs to be autoregressive so it cannot look at the whole sequence,
3. So, how can it look at the whole sequence and not look at the whole sequence at the same time?

I was also confused about this, so I'll try to explain.

Say we want to train on a sequence of tokens, i.e. s = [1, 13, 64, 40]. The next token to predict after 40 is 50. Now, instead of having one training sample, s, we can have several sequences from s, such as s1 = [1] (with target 13), s2 = [1, 13] (with target 64) etc., giving a total of 4 sequences.

In total we have the following training data

sample = [
 [1],
 [1, 13],
 [1, 13, 64],
 [1, 13, 64, 40]]

targets = [13, 64, 40, 50]

As usual, Tensors do not like to have weird shapes. Therefore, the training data will look something like this

sample = [
 [1, -inf, -inf, -inf],
 [1,  13,  -inf, -inf],
 [1,  13,   64,  -inf],
 [1,  13,   64,   40]]

In this new tensor, the -inf is the masking! Now you might say, "that's just sequence padding", and it is! The only reason that they call it masking instead of padding is because the architecture in the paper has an encoder (next to the decoder), in which the complete sequences are used without masking. However, when only using the decoder (as in LLMs), masking is essentially just padding.

Answer 2

During training in a transformer model, the entire target sequence is indeed fed into the decoder as input. However, the output sequence is still generated autoregressively, one token at a time. This means that the decoder generates each token based on the previously generated tokens, rather than generating all the tokens in parallel.

To achieve this, the decoder uses a masked self-attention mechanism, where the attention weights are masked to prevent attending to future positions in the sequence. This ensures that the decoder only generates tokens based on the previously generated tokens, as it would during inference.

The reason for feeding the entire target sequence as input during training is to enable parallel computation of the decoder layers. By doing so, the model can process multiple tokens in parallel, which speeds up training compared to generating the tokens one at a time. However, the autoregressive nature of the decoder output generation is still maintained, as this is crucial for ensuring that the model can generalise to unseen sequences during inference.

Attention is all you need, p. 3:
"We also modify the self-attention sub-layer in the decoder stack to prevent positions from attending to subsequent positions. This masking, combined with fact that the output embeddings are offset by one position, ensures that the predictions for position i can depend only on the known outputs at positions less than $i$."