# Last linear layer of the decoder of a transformer

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.

• Anyway I am starting to be quite convinced that I am right. If the final linear layer were processing all the vectors output of a decoder than the number of parameters of this dense layer would be dependent on the sequence length, which would give problems in handling inputs with different lengths (which we need to handle for language translation) Commented Aug 10, 2022 at 12:35

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.

### Edit

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

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.

• Thanks a lot for your answer. Give me please some time to look at the implementation... In the meanwhile there are some clarifications I would need to better understand. Commented Aug 11, 2022 at 6:23
• (1) You say "the inputs to the encoder and the decoder the next time you apply the model to the sequence (to get the next token wT+1) is going to be (w1,w2,...,wT)". --> what setting do you have in mind ? I am thinking in this thread about sentence translation. For example I would like to translate in german "The dog is on the table". For my understanding the encoder is used (forward pass) just once with the whole english sentence "The dog is on the table". Instead the decoder, as you say, receives as input the translated sentence. "Der" , than "Der Hund", than "Der Hund ist", and so on... Commented Aug 11, 2022 at 6:25
• (2) "The output of the decoder (after sampling) will be (w2,w3,...,wT+1)."--> Since we already have in your example "(w1,w2,...,wT)". Let's call "(w'2,w'3,...,w'T+1)" the output of the decoder, since there is no guarantee that w2=w2',... and so on... The only word we will use in the following steps is w'T+1, right ? (see also next question) Commented Aug 11, 2022 at 6:31
• (3) Then, applying the model again to get wT+2, the new input to the encoder and the decoder is going to be (w2,w3,...,wT,wT+1) --> According to the last question, I thought that the new input would have been (w1,w2,w3,...,wT,w'T+1) (see also video in section "The Decoder Side" of the link "The illustrated transformer) Commented Aug 11, 2022 at 6:32
• Let me update my question for the translation use case. Commented Aug 11, 2022 at 19:43

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).