Softmax & Probability Distributions

Welcome to our next lesson. In our previous session, we examined the two primary pre-training objectives for Transformers: Causal Language Modeling (CLM) for generation and Masked Language Modeling (MLM) for understanding. We established that regardless of the specific objective, the model's ultimate task is to predict a token from its vocabulary based on some context.

This lesson drills down into the final, crucial step of that process. We will explore how a Transformer takes its internal, high-dimensional vector representations and converts them into a meaningful probability distribution over the entire vocabulary. Your goal is to understand the mechanics and purpose of the two key components that achieve this: a final linear layer and the softmax function. This is the bridge between the abstract, learned representations inside the model and the concrete, probabilistic output we use for prediction.

From Hidden States to Vocabulary Scores

Throughout the Transformer stack, information is processed and transformed into a sequence of vectors. The output of the final Transformer block is a sequence of hidden state vectors, each with a dimension of $d_{model}$ . For the task of predicting a specific token (either the next one in CLM or a masked one in MLM), we focus on the single hidden state vector corresponding to the position we want to predict.

This vector, let's call it $h_{L} \in R^{d_{model}}$ , is rich with contextual information. However, it exists in an abstract high-dimensional space. Our goal is to map this vector to a score for every single token in our model's vocabulary, which has a size of $V$ . This is where the Language Model Head comes in.

This diagram illustrates the role of the Language Model Head. It takes the final hidden state vector from the Transformer, $h_N^L$, and processes it through two stages to produce a probability distribution over the vocabulary.

The Language Model Head is typically composed of two parts: a linear transformation and a softmax activation function. Let's examine each one.

For a clear, formal overview of this entire process, the following reading from Stanford's CS124 course is excellent.

[PDF] Introduction to Large Language Models - Stanford University

This section of the Stanford lecture slides provides a concise yet comprehensive explanation of the language modeling head.

Please read the part of the slide deck that discusses the Language Modeling Head. Start from the heading How do we turn it into... and continue through the explanation of the unembedding layer and softmax. Pay close attention to the diagrams showing the change in vector dimensions from d to |V|.

Step 1: The Final Linear Layer and Logits

As the reading explained, the first step is to project the $d_{model}$ -dimensional hidden state vector into a $V$ -dimensional vector. This is achieved with a simple linear layer.

Since you're familiar with linear algebra, you'll recognize this as a matrix multiplication. The hidden state vector $h_{L}$ is multiplied by a weight matrix, often called the "unembedding matrix" $W_{U} \in R^{d_{model} \times V}$ .

u = h_{L} W_{U}

The resulting vector $u \in R^{V}$ is a vector of raw, unnormalized scores. Each of the $V$ components in this vector corresponds to a token in the vocabulary. These scores are called logits. A higher logit for a particular token suggests that the model thinks that token is more likely to be the correct one, but these are not yet probabilities. They can be any real number, positive or negative.

The term "unembedding" is used because this operation is conceptually the reverse of the input embedding layer, which maps a token ID (from vocabulary space) to a $d_{model}$ -dimensional vector. This final linear layer maps a $d_{model}$ -dimensional vector back to the vocabulary space. In many models, the weights of this unembedding matrix are even "tied" to the weights of the input embedding matrix (specifically, $W_{U} = W_{E}^{T}$ ), which saves parameters and has been found to work well in practice.

The following video from 3Blue1Brown explains this "unembedding" step and sets the stage for what comes next.

Transformers, the tech behind LLMs | Deep Learning Chapter 5

This video provides a clear, intuitive explanation of the final layers of a Transformer.

Watch the segment from this point. The narrator explains how the final hidden state vector is mapped to a list of scores for each token in the vocabulary using the "unembedding matrix." This is the linear layer that produces the logits.

For a more rigorous mathematical definition and a practical implementation note, the following text is very helpful.

Final Linear Layer and Softmax Output - ApX Machine Learning

This resource from ApX Machine Learning provides the mathematical formulation for the linear transformation.

Read the introduction and the section on The Final Linear Transformation. Note the explicit formula for calculating the Logits tensor and the implementation note about weight sharing.

Step 2: The Softmax Function

The logits are raw scores, not probabilities. To convert the logit vector $u$ into a valid probability distribution, we apply the softmax function. The softmax function takes a vector of $V$ real numbers and transforms it into a vector of $V$ real numbers that collectively sum to 1, with each element being between 0 and 1.

For each logit $u_{i}$ in the vector $u$ , its corresponding probability $p_{i}$ is calculated as:

p_{i} = softmax (u_{i}) = \frac{e ^{u_{i}}}{\sum _{j = 1}^{V} e ^{u_{j}}}

This formula does two things:

Exponentiation ( $e^{u_{i}}$ ): This makes all the scores positive. It also has the effect of amplifying differences; a small increase in a large logit results in a much larger increase in its exponentiated value compared to the same increase in a small logit.
Normalization: Dividing by the sum of all exponentiated scores ensures that the entire output vector sums to 1, creating a valid probability distribution.

The name "softmax" is fitting because it acts as a "soft" version of the max function. Instead of outputting a 1 for the maximum value and 0s for all others, it assigns a high probability to the token with the highest logit and smaller, but non-zero, probabilities to others.

This diagram shows the complete flow. The last hidden state from the Transformer is passed to a linear layer to produce logits. The softmax function then converts these logits into a probability distribution, from which the most probable next token can be identified.

The 3Blue1Brown video you just watched has an excellent follow-up segment that provides one of the best intuitions for the softmax function and introduces a fascinating parameter called "temperature".

Transformers, the tech behind LLMs | Deep Learning Chapter 5

This segment continues from the last one, providing a deep dive into the softmax function.

Please watch from this point. Focus on how softmax turns an arbitrary list of numbers (logits) into a valid probability distribution. Pay special attention to the explanation of the temperature parameter, which controls the "peakedness" of the distribution. This concept is crucial for controlling the randomness of text generation, a topic we will revisit in a later module.

As you can see, this two-step process—a linear projection followed by a softmax normalization—is a simple yet powerful mechanism that allows the model to make a final, interpretable prediction from its complex internal state.

Finally, to see how this is represented in code, the following clip from Andrej Karpathy's "Let's build GPT" series is illuminating. He explicitly adds a "language modeling head" (lm_head) which is simply a PyTorch nn.Linear layer. Your background in CS and development will make this connection between theory and implementation immediately clear.

Let's build GPT: from scratch, in code, spelled out.

This video by Andrej Karpathy demonstrates building a GPT model from scratch in PyTorch.

Watch the short segment from this point. He defines self.lm_head as a linear layer that maps from the embedding dimension (n_embd) to the vocabulary size (vocab_size). This is the practical implementation of the unembedding matrix we've been discussing.

Conclusion

In this lesson, we have demystified the final output stage of a Transformer model. You now understand the precise mechanism that translates the model's learned knowledge into a concrete prediction.

Here are the key takeaways:

From Hidden State to Logits: The final hidden state vector from the Transformer, which is rich in context, is first projected into a vector the size of the vocabulary using a linear layer. This layer is often called the language model head or unembedding layer.
Logits as Raw Scores: The output of this linear layer is a vector of raw, unnormalized scores called logits. Each element in the logit vector corresponds to a token in the vocabulary.
From Logits to Probabilities: The softmax function is applied to the logit vector to convert it into a valid probability distribution. It does this by exponentiating the logits (making them positive) and normalizing them so they sum to 1.
The Final Output: The result is a probability for every single token in the vocabulary, representing the model's belief about which token should appear at the predicted position.

We now have a probability distribution. But during training, how does the model know if its predicted distribution is "good" or "bad"? How do we quantify the error and use it to update the model's billions of parameters? That is the subject of our next lesson, where we will derive and explore the cross-entropy loss function.

Recommended resources

From Hidden States to Vocabulary Scores

Step 1: The Final Linear Layer and Logits

Step 2: The Softmax Function

Conclusion