# How is the next token predicted in transformers?

In the transformer (or GPT/decoder only), at the end of the decoder blocks but before the final linear layer you have X vectors (for the X tokens at the input of the decoder). We then want to compute the probabilities for the next token of the sequence - what do we then feed to the linear layer? Is it the last embedding corresponding to the hidden state of the last token in the input sequence?

I've seen some tutorials on youtube on how to make mini gpts but I never quite understood why they feed the entire X vectors/hidden states at the end of the decoder blocks to the linear layer and not just the last vector/hidden state... Wouldn't you have X probability distributions when in reality you only want one? And if we do want the X probability distributions then wouldn't we be completely missing the point of the masked self attention since we would be trying to predict words that are already in the input sequence, so essentially "cheating"?

• Maybe this video explains it for you. It's about the Attention is all you need paper. Commented Apr 21, 2023 at 8:04
• That video never talks about how the outputs are fed into the linear layer, but thank you Commented Apr 23, 2023 at 17:35

Welcome to AI stack exchange!

I understand the confusion. Inference (next token prediction) seems really counterintuitive and inefficient for transformers. And it is! The transformer is very efficient during training because it can be parallelized. It is, however, inefficient at inference because it cannot be parallelized.

For transformer inference, you feed the context (your prompt) to the transformer model. It predicts the next word for each of the words in the prompt, but you only need the prediction for the last one.

A bit of pseudocode might help in understanding how a transformer can be used to generate new tokens:

# Start with some context of tokens
context = ...

# Generate new tokens
for i in range(N_TOKENS_TO_GENERATE):
prediction = transformer(context)                # Get predictions for context
next_token = multinomial(prediction.get_last())  # Sample from multinomial distribution
context = concatenate((context, next_token))     # Create new context


Now, this is the intuitive way of doing it. There are most likely tons of small things you can do to optimize all of the stuff and make inference more efficient. However, you cannot get around having to feed the context in every time you add a new word/token. This is also why an application such as ChatGPT is generating stuff word for word.

A small note on the side: you talk about 'hidden-states' in the transformer, as if there is a recurrence going on (such as in GRUs/LSTMs/RNNs). However, transformers have no such recurrence and hidden-states and operate solely using the concept of attention (hence the paper's title 'attention is all you need', alluding to the fact that they don't use recurrence).

Hope this helps :)

• Thank you! What about in training? What is the point of feeding the context embeddings to the linear layer and not only the last linear layer, like it's done in github.com/karpathy/nanoGPT? Also, some of the papers I read/people I talk to sometimes talk in hidden states in transformers (maybe out of habit because of RNNs), so that is why I also mentioned hidden states ahah Commented Apr 21, 2023 at 21:37
• im not sure i understand your first followup question. Could you elaborate a little bit on what you mean? Regarding the hidden-state, its very understandable but there is a notable difference, because in transformers you dont have a 'state' after a prediction which you can reuse. You have to recompute everything for the new context. Commented Apr 22, 2023 at 8:10
• So what I meant was this: why are all the context embeddings at the end of all decoder blocks fed into the linear layer and not only the last context embedding at the end of the decoder blocks? Wouldn't it be better to only use a single vector at the linear layer instead of the whole context vectors? In other words, in the original gpt paper, why is it that the matrix h_n is fed to the last linear layer and not only the last token embedding (h_n[-1])? Commented Apr 23, 2023 at 14:26
• Because everything together represents the embedding, not just the last bit. Its a feedforward network, not a recurrent network. Commented Apr 23, 2023 at 17:23
• But how does that even work? Btw I'm referring to equation 2 in that paper. If you multiply a matrix by a matrix you get a matrix... How do you get a probability distribution over all tokens for the next token in the sequence? Commented Apr 23, 2023 at 17:31

No - each next-token prediction comes from a single one of the output vectors as you suspected. It has to, because otherwise there is no way to parallelize the predictions during training using a consistent set of parameters.

My understanding is from https://transformer-circuits.pub/2021/framework/index.html, section "High Level Architecture":

What you are referring to as "X vectors", they refer to as T(t) in that figure, of shape [n_context, n_vocab], which are logits. The input to the decoder, t is the [n_context, n_vocab] shaped tensor of one-hot encoded tokens (see the "Notation" section of that article at the end).

Note also that the original "Attention is All You Need" paper re-uses the embedding matrix as a transpose to de-embed each of the X vectors coming out of the final multi-head attention decoder layer. (See section 3.4 in the original paper)