Calculating Cross-Entropy Loss for Next-Token Prediction

Introduction

In our last lesson, we established how the softmax function transforms a language model's raw output scores (logits) into a clean probability distribution over the entire vocabulary. For a given sequence, the model can now tell us, for example, that the word "mat" has a 75% probability of being the next token, "rug" has 15%, and so on.

This is a crucial step, but it doesn't answer the most important question for training: How "good" is this prediction? To improve the model through training, we need a way to quantify its error. If the actual next word was "mat", a 75% prediction is good, but not perfect. If the actual word was "dog", it's a very poor prediction.

This lesson introduces the cross-entropy loss function, the industry-standard method for measuring this error. We will learn how to compute it and understand why it is so effective for training language models.

By the end of this lesson, you will be able to:

• Define the purpose of a loss function in model training.
• Represent the "true" next token as a one-hot encoded vector.
• Compute the cross-entropy loss for a single next-token prediction from the model's probabilities.
• Understand why cross-entropy is particularly effective for training classification models.
• Implement cross-entropy loss in PyTorch.

1. Defining the Prediction Error

Our goal is to compare two things:

1. The Model's Prediction: A probability distribution across the entire vocabulary (e.g., [0.15, 0.06, 0.76, ...]).
2. The Ground Truth: The single, correct next token from our training data.

To compare them mathematically, we first need to express the ground truth in the same format as the prediction—as a probability distribution. Since we know the correct next token with 100% certainty, we represent it using a one-hot vector. This is a vector that is zero everywhere except for a single 1 at the index corresponding to the correct token.

For a vocabulary of [the, cat, sat, on, mat]:

• If the model sees "The cat sat on", the target token is "the".
• The one-hot vector for "the" would be [1, 0, 0, 0, 0].

The training process for an autoregressive model involves making a prediction at each position in a sequence. The following resource clarifies how we generate these prediction-target pairs from a single sentence.

As the article shows, for each token in an input sequence, the model predicts the next token. The loss is then calculated by comparing the model's predicted distribution to the one-hot vector of the actual next token.

2. Cross-Entropy Loss: The Core Concept

Cross-entropy loss measures the dissimilarity between two probability distributions. In our case, it measures how different the model's predicted distribution is from the one-hot ground truth distribution. The goal of training is to adjust the model's parameters to minimize this loss, effectively pushing the predicted distribution to be as close as possible to the true one.

This beautiful visualization illustrates the objective. We want to increase the probability assigned to the correct word while decreasing the probabilities of all incorrect words.

The following video provides a fantastic, intuitive explanation of cross-entropy.

The Simplified Formula

As the video and resources explain, the general formula for cross-entropy is:
$L_{CE}=-{\sum }_{j=1}^V{y}_j\log ({\hat{y}}_j)$
Where:

• $V$ is the vocabulary size.
• $y_j$ is the true probability of the $j$-th token (from the one-hot vector).
• ${\hat{y}}_j$ is the predicted probability of the $j$-th token (from the softmax output).

Because the true distribution $y$ is a one-hot vector, $y_j$ is 0 for all incorrect tokens and 1 for the single correct token. This causes all terms in the summation to become zero except for the one corresponding to the correct token.

This simplifies the entire calculation to a beautifully simple formula:

$L_{CE}=-1\cdot \log ({\hat{y}}_{\text{correct}})=-\log ({\hat{y}}_{\text{correct}})$

The cross-entropy loss for a single prediction is simply the negative logarithm of the probability the model assigned to the correct token.

3. Step-by-Step Computation

Let's compute the loss for a prediction.
Suppose after seeing "The cat sat on the", our model produces a probability of 0.75 for the correct next token, "mat".

• Loss Calculation:
  $L=-\log (0.75)\approx -(-0.287)=0.287$

Now, consider a poor prediction where the model only assigned a probability of 0.01 to "mat".

• Loss Calculation (Bad Prediction):
  $L=-\log (0.01)\approx -(-4.605)=4.605$

Notice two key properties:

1. Loss is always non-negative. Since probabilities $\hat{y}$ are between 0 and 1, their logarithm is always negative or zero. The negative sign in the formula flips this to a positive loss.
2. Lower probability means higher loss. As the predicted probability of the correct token approaches 0, the loss approaches infinity. As it approaches 1 (a perfect prediction), the loss approaches 0.

Why is This So Effective?

This behavior is exactly what we want from a loss function. It heavily penalizes predictions that are confidently wrong, which in turn creates a large gradient, guiding the model to make significant corrections during training.

This next video segment explains this powerful property by comparing cross-entropy to a more intuitive alternative, the sum of squared residuals.

Loss Over a Full Sequence

In practice, we calculate this loss for every token prediction across a batch of sequences. The final loss for the entire batch is simply the average of all the individual token losses.

4. Practical Implementation in PyTorch

Your familiarity with PyTorch and CNNs is a great asset here. The concept is identical to multi-class classification in computer vision. However, there is a critical implementation detail you must know.

While the mathematical formula for cross-entropy loss requires probabilities (the output of softmax), PyTorch's torch.nn.CrossEntropyLoss function is optimized for performance and numerical stability. It expects raw logits as input, not probabilities.

There are two key takeaways when using torch.nn.CrossEntropyLoss:

1. Input: It takes the model's raw logits directly. It performs the softmax and log calculation internally.
2. Target: It expects the correct labels as a tensor of integer indices, not one-hot encoded vectors.

Let's see this with a concrete example.

import torch
import torch.nn as nn




# --- Setup ---
# Imagine a small vocabulary of 5 words.
# The model has produced raw logits for a single next-token prediction.
# Logits: [cat, dog, mat, sat, the]
logits = torch.tensor([[1.2, 0.1, 3.8, 2.1, -1.5]]) # Shape: [batch_size, num_classes]




# The correct next token is "mat", which is at index 2 in our vocabulary.
targets = torch.tensor([2]) # Shape: [batch_size]




# --- Loss Calculation ---
# Initialize the loss function
loss_function = nn.CrossEntropyLoss()




# Compute the loss
loss = loss_function(logits, targets)

print(f"Logits: {logits}")
print(f"Target index: {targets}")
print(f"Cross-Entropy Loss: {loss.item():.4f}")




# --- Manual Verification (to connect with the math) ---
# 1. Apply softmax to logits
probs = torch.softmax(logits, dim=1)
print(f"Probabilities (from softmax): {probs}")




# 2. Get the probability of the correct class (index 2)
prob_of_correct_class = probs[0, 2]
print(f"Probability of 'mat': {prob_of_correct_class.item():.4f}")




# 3. Calculate negative log of that probability
manual_loss = -torch.log(prob_of_correct_class)
print(f"Manual Loss (-log(p)): {manual_loss.item():.4f}")

Running this code would show that loss and manual_loss are identical, confirming that nn.CrossEntropyLoss is indeed computing the negative log probability of the correct class from the raw logits.

Conclusion

You've now learned about the second of the two fundamental components of language model training. With softmax to create a probability distribution and cross-entropy loss to measure the error of that distribution, we have everything we need to guide the model's learning process.

Key Takeaways:

• Cross-entropy loss measures the dissimilarity between a model's predicted probability distribution and the true, one-hot distribution of the target token.
• For next-token prediction, this simplifies to calculating the negative logarithm of the probability assigned to the correct token.
• A lower probability for the correct token results in a higher loss, providing a strong training signal.
• PyTorch's nn.CrossEntropyLoss is the standard, efficient implementation. It requires raw logits for the prediction and integer class indices for the target.

Preview of the Next Lesson:

We now know how to calculate the error for a single predictive step. But how do we manage this process over an entire sequence of text? When the model predicts the third token, should it be given the actual second token as input, or the token it predicted for the second position? The next lesson will cover teacher forcing, the standard training procedure for autoregressive models that answers this exact question.