LLM output phase after transformers
When an LLM runs inference, you can think of three main phases:
- Preparation: tokenization, embeddings, and positional encodings → tensor X
- Transformer blocks: multi-head attention + feed-forward → hidden states H
- Output (next-token): turn the final hidden states into an output token.
This blurb focuses on the output phase.
Definitions
- V: vocabulary size (number of possible tokens).
- B: batch size (number of sequences processed in parallel).
- T: sequence length (tokens per sequence).
- D: d_model, hidden dimension for embeddings and transformer states.
- X: token embeddings with positional encoding, shape (B, T, D).
- H: final hidden states from the transformer stack, shape (B, T, D).
- Logits: raw, unnormalized scores over the vocabulary; one logit per token in V.
- Softmax: converts a vector of logits into a probability distribution: outputs are in [0, 1] and sum to 1.
Decoding / probability selection methods
Given a probability distribution over the vocabulary for the next token, common selection strategies are:
- Greedy decoding: always pick the single highest-probability token (fully deterministic).
- Top-k sampling: restrict to the k most probable tokens, renormalize their probabilities, then sample from that set (controlled randomness).
- Top-p sampling (nucleus sampling): sort tokens by probability and include the smallest set whose cumulative probability ≥ p; renormalize and sample from that set.
Output phase data flow
For each step of next-token prediction:
- The transformer stack produces hidden states H with shape (B, T, D).
- A final linear projection maps H to vocabulary space, producing logits with shape (B, T, V).
- Softmax converts logits into probabilities over the vocabulary.
- A decoding strategy (greedy, top-k, top-p, etc.) selects one token based on those probabilities.
- The selected token is decoded back to text and appended to the sequence. The process repeats for the next position until it encounters an 'end-of-sequence' token or max length.
- The final sequence is returned as the model's output
References
More details about the various steps:
- jalammar.github.io/hidden-states/
- machinelearningmastery.com/how-llms-choose-their-words-a-practical-walk-through-of-logits-softmax-and-sampling/
- www.kenmuse.com/blog/how-temp-topk-topp-minp-control-llm-output/
- mikexcohen.substack.com/p/llm-breakdown-46-transformer-outputs
- huggingface.co/blog/mlabonne/decoding-strategies
- www.assemblyai.com/blog/decoding-strategies-how-llms-choose-the-next-word
- python.plainenglish.io/understanding-decoding-strategies-greedy-beam-search-top-k-top-p-27fbefcc3295
- opt-ml.org/papers/2025/paper160.pdf
- en.wikipedia.org/wiki/Top-p_sampling
- huyenchip.com/2024/01/16/sampling.html
- www.aussieai.com/research/top-k-decoding
- www.youtube.com/watch?v=JqgsxC4guK4
- www.intoai.pub/p/decoding-strategies-in-llms
- www.linkedin.com/posts/baijayantaroy_genai-llm-decoding-strategies-explained-activity-7316111234087464961-0jl6
- huggingface.co/docs/transformers/en/main_classes/output
- mbrenndoerfer.com/writing/logit-lens
- github.com/huggingface/transformers/issues/15392
NOTES
- This is a simplified explanation of the process. LLM research is coming up with new stuff daily.
- This article was produced from my detailed description, reviewed and enhanced by Perplexity