# How do the linear + softmax layers give out word probabilities in Transformer network?

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?

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.

• 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.
– k-c
Jul 8, 2021 at 18:34
• 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.
– noe
Jul 8, 2021 at 18:37