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

Question

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?

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.