Last linear layer of the decoder of a transformer

Question

I am learning the transformers architecture from these two sources:

https://arxiv.org/pdf/1706.03762.pdf

https://jalammar.github.io/illustrated-transformer/

I just wanted to ask about the final step in the decoder. Let's fix testing time. As I understand, the decoder starts with an input of dimension $(N_{words},d_{emb})$, where $N_{words}$ is the number of words already predicted and $d_{emb}$ is the embedding dimension.

Now if we "follow" the following decoder steps, at each step (after e.g. the attention layers) we should have a vector of dimension $(N_{words},d_{model})$ where $d_{model}$ is the model dimension. In other words, up to the final linear layer we have $N_{words}$ vectors which are $d_{model}$-dimensional.

Are all these $N_{words}$ vectors fed into the last linear layer (before the softmax) or, as I suspect, only the last of these vectors is used ? In the latter case the last linear layer would be a matrix of dimension $d_{model}\times N_{vocab}$, where $N_{vocab}$ is the vocabulary dimension.

Is this correct ? Are there any issues in what I wrote ? Unluckily from the online sources I was not able to clarify this point...

PS: I conjectured that the last linear layer is using just the last vector, because than I would understand what happens in training time, one would just use in that case all the output vectors from the decoder, instead of just the last one, to have a parallelized prediction.

Answer 1

I agree with this:

PS: I conjectured that the last linear layer is using just the last vector, because than I would understand what happens in training time, one would just use in that case all the output vectors from the decoder, instead of just the last one, to have a parallelized prediction.

Indeed when you look at the code of Tensorflow shared previusly you see this line

# select the last token from the seq_len dimension
  predictions = predictions[:, -1:, :]  # (batch_size, 1, vocab_size)

where they take the last token of the sequence as a prediction for the next token. And this is consistent with the fact that the input sequence to the decoder is shifted to the right by adding a <START> token at the beginning (see transformers original paper section 3.1-Decoder).

So to correct the previous answer, the input and output sequence of the decoder have the same length and the last vector is used for the next token prediction which is then added at the end of the input sequence.

Answer 2

Edit

Based on the comments to the original version of this answer, OP indicated that the use case was translation between two languages.

Answer:

At sampling time, the last linear layer of the decoder is going to output a sequence whose length is incremented by one each time you apply the encoder-decoder transformer to the input sentence.

Let's take a practical example, with $w$ denoting the words in the original sentence and $w^{i}$ those in the target language after iteration $i$ of applying the transformer model.

If you have already sampled a sequence $(w^1_1, w^2_2, ..., w^{l}_l)$, the inputs to the encoder and the decoder the next time you apply the model (to get the next token $w^{l+1}_{l+1}$) are going to be respectively $(w_1, w_2, ..., w_T)$ (the original sentence) and $(w^1_1, w^2_2, ..., w^l_l)$. The output of the decoder (after sampling) will be $(w^{l+1}_1, w^{l+1}_2, ..., w^{l+1}_l, w^{l+1}_{l+1})$, but we're only going to keep $w^{l+1}_{l+1}$.

Then, applying the model again to get $w^{l+2}_{l+2}$, the new inputs to the encoder and the decoder are going to be $(w_1, w_2, ..., w_T)$ (again the original sentence) and $(w^1_1, w^2_2, ..., w^l_l, w^{l+1}_{l+1})$. The output of the decoder will be $(w^{l+2}_1, w^{l+2}_2, ..., w^{l+2}_l, w^{l+2}_{l+1}, w^{l+2}_{l+2})$, and we'll keep only $w^{l+2}_{l+2}$.

This continues until the sentence is fully translated.

At each step, we give the whole original sentence to the encoder so that the model can build its translation for the next word by looking at the whole sentence. And we give the sentence translated so far to the decoder so that the translation of the next word can attend to what has already been translated.

To confirm, you can have a look at what they do in the TensorFlow transformer tutorial.

Answer 3

Well well, this text in the link posted in one of the answers above (https://www.tensorflow.org/text/tutorials/transformer) gives the answer to the question.

Note: The model is optimized for efficient training and makes a next-token prediction for each token in the output simultaneously. This is redundant during inference, and only the last prediction is used. This model can be made more efficient for inference if you only calculate the last prediction when running in inference mode (training=False).