I have followed this tutorial https://www.youtube.com/watch?v=U0s0f995w14 to create a minified version of a transformer architecture but I am confused about the final shape of the output.

Heres the code: https://github.com/aladdinpersson/Machine-Learning-Collection/blob/558557c7989f0b10fee6e8d8f953d7269ae43d4f/ML/Pytorch/more_advanced/transformer_from_scratch/transformer_from_scratch.py#L2

on the final lines (slightly modified):

print(trg[:, :-1].shape)
out = model(x, trg[:, :-1])

The output shapes don't seem to make sense

torch.Size([2, 9]) #input sentence (num_examples, num_tokens)
torch.Size([2, 7]) #generated sentence so far (num_examples, num_tokens_generated)
torch.Size([2, 7, 10]) # probabilities for next token (num_examples, ???, size_vocab)

The transformer is supposed to predict the next token across 2 training examples (which is why theres a two for the number of examples and a 10 for the size of the vocab), by generating probabilites for each token in the vocab. But I can't make sense of why theres a 7 there. The only explanation I can come up with is that it outputs all predictions simulatenously, but that would require feeding the outputs iteratively through the transformer, but that never happens (see lines 267-270).

So is there a mistake or am I not understanding something correctly? What is that output shape supposed to represent?

Can somebody make sense of this?


1 Answer 1


7 is the length of the target sequence passed as argument to the model, which is trg[:, :-1], that is, the target sequence except the last token. The last token is removed because it contains either the end-of-sequence token of the longest sentence in the batch or padding tokens of the shorter sequences in the batch, and therefore it is useless.

The output of the decoder is of the same length as its input. The shape of trg[:, :-1] is [2, 7], so the shape of the output is the same.

Note that in the video they are invoking the model in an unusual way, because they are passing a whole target sequence to the model but they are not training it. Normally, the model would be used in one of the following ways:

  • In training mode, receives a full target sequence and its output is used to compute the loss and update the network weights via gradient descent.
  • In inference mode, is used auto-regressively, that is, we decode token by token, incorporating each new predicted token into the input for the next step.

I guess they used the model this way just to illustrate that the model works.

  • $\begingroup$ I'm confused. According to attention is all you need paper, the output is a set of probabilities over the vocab for the next token. Why would it return the same sentence it got fed into it? $\endgroup$ Commented Jun 11, 2023 at 17:47
  • $\begingroup$ The output is the probability distribution over the vocab for all the tokens. At inference time, we choose to ignore everything but the last position to make the prediction of the next token. At training time, we take the output at all the positions to compute the loss. $\endgroup$
    – noe
    Commented Jun 11, 2023 at 17:49
  • $\begingroup$ Okay, but then shouldn't the output be 8 then? Or is 7 the context size and in the real application the rest is simply padded with 0s? Also wouldn't the Transformer simply learn a 1-1 mapping of everything but the the last token? $\endgroup$ Commented Jun 11, 2023 at 18:13
  • $\begingroup$ The last position is usually removed. At inference time, this removal is optional, you may or may not remove it; removing the last token does not hurt because it either contains end-of-sequence or padding, neither of which is not needed in any case. In training, the removal of the last position is actually needed because we need to shift the expected outputs one position (so that the expected output at the first position is the target input at the second position), which means that we lose one token in the expected output and, therefore, we need to remove the last token of the target input. $\endgroup$
    – noe
    Commented Jun 11, 2023 at 18:20
  • 1
    $\begingroup$ If you have further questions, please create a new question on the site. $\endgroup$
    – noe
    Commented Jun 12, 2023 at 9:06

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.