I am trying to implement a transformer network from scratch in pytorch to understand it. I am using The illustrated transformer for guidance. The part where I am stuck is about how do we go from the output of the final decoder layer to linear + softmax.

From what I have understood, if we have a batch of size B, max output seq length M, embedding dimension D, and vocab size V, then the output of the last decoder layer would be BxMxD which we have to turn into a vector of probabilities of size BxV so that we can apply softmax and get next predicted word. But how do we go from a variable size MxD matrix to a fixed-length V vector?

This post says we apply the linear layer to all M vectors sequentially:

That's the thing. It isn't flattened into a single vector. The linear transformation is applied to all M vectors in the sequence individually. These vectors have a fixed dimension, which is why it works.

But how do we coalesce those transformed vectors into just one single vector? Do we sum them up?


1 Answer 1


Your understanding is not correct.

You go from a $B \times M \times D$ tensor to a $B \times M \times V$ tensor (i.e. the logits). As you can see, in the final tensor we have M vectors of dimension V (one vector per token), not just a single vector.

To obtain the $B \times M \times V$, you just perform a matrix multiplication.

This applies to Transformers, but also to most sequence generation models, like LSTMs.

  • $\begingroup$ So the decoding happens in just one pass? What I have seen is that the encoding happens in one go, then we sequentially generate the output sentence one word at a time, each time giving the sentence generated up to that point as the input. $\endgroup$
    – k-c
    Commented Jul 8, 2021 at 18:34
  • $\begingroup$ At training time, everything is predicted in one pass. At inference time, the predictions are generated token by token. You can check this answer for details on this. $\endgroup$
    – noe
    Commented Jul 8, 2021 at 18:37

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.