Cross-Entropy Loss for Next-Token Prediction

Next-token prediction is classification over the vocabulary

Before any loss makes sense, be precise about what the model outputs. At each position, a transformer produces a vector of logits z ∈ R^V, one real number per token in a vocabulary of size V (often 32k–128k). A softmax turns those logits into a probability distribution p = softmax(z) over the whole vocabulary: the model’s guess for ‘what token comes next.’

So next-token prediction is just multi-class classification with an enormous number of classes, repeated once per position in the sequence. The ‘label’ at each position is the single token that actually appears next in the training text — a hard, known answer. That is the crucial structural fact the loss exploits: we are not comparing two soft distributions in general, we are comparing the model’s distribution against a target that puts all its mass on one correct token. Everything that follows — why the loss reduces to a single log, why the gradient is so clean — flows from that one-hot target. Shapes to keep in mind: logits z: [B, T, V] for batch B and sequence length T, targets y: [B, T] as integer token ids.

Maximum likelihood: the goal before the loss

Training a language model starts as a maximum likelihood problem, not a loss-minimization one. We have a corpus of real text, and a model with parameters θ that assigns a probability to any sequence. The principle is simple: choose θ to make the observed text as probable as possible. Autoregressive factorization lets us write the probability of a sequence as a product of per-token conditionals:

p_θ(x_1, ..., x_T) = ∏_t p_θ(x_t | x_1, ..., x_{t-1})

Maximum likelihood says: maximize this product over the whole dataset. But a product of thousands of probabilities, each below 1, is a numerical disaster — it underflows to zero almost instantly, and its gradient is a nightmare of chained product rules. The fix is the single most important move in the derivation: take the logarithm. Because log is monotonically increasing, whatever θ maximizes the product also maximizes its log — we lose nothing by optimizing the log-likelihood instead, and we gain a sum, numerical stability, and additive gradients. That one substitution is where cross-entropy is born.

From likelihood to negative log-likelihood

Applying log to the product turns it into a sum, because log(ab) = log a + log b:

log p_θ(x_1..x_T) = Σ_t log p_θ(x_t | x_<t)

Maximizing this log-likelihood is identical to minimizing its negative — and optimizers are conventionally written as minimizers. So we define the negative log-likelihood (NLL) as the loss:

NLL(θ) = - Σ_t log p_θ(x_t | x_<t)

Each term -log p_θ(x_t | x_<t) is the loss contributed by one position: the negative log-probability the model assigned to the token that actually came next. If the model was confident and correct (p → 1), -log p → 0 — no penalty. If it was confident and wrong (p → 0), -log p → ∞ — an unbounded penalty. That asymmetry is deliberate and important: NLL punishes confident mistakes savagely and rewards calibrated confidence, which is exactly the behavior we want from a probabilistic predictor. This per-token NLL is the cross-entropy loss; the next section shows why the two names describe the same number.

Cross-entropy: the general definition

Cross-entropy is a quantity defined between two probability distributions over the same set of outcomes: a target distribution y and a predicted distribution p. It measures the average number of nats (natural-log units) needed to encode samples from y using a code optimized for p:

H(y, p) = - Σ_{i=1}^{V} y_i log p_i

It is minimized, for a fixed y, exactly when p = y: the model matches the target distribution. That is what makes it a sensible loss — driving H(y, p) down drives the prediction toward the target. In general both y and p can be soft, but in language modeling the target is special. The correct next token is known, so y is a one-hot vector: y_c = 1 for the true token index c, and y_i = 0 for every other i. Substituting a one-hot y into the sum annihilates every term except the one at i = c, and the whole expression collapses to a single log — the connection we make explicit next. This is why the loss you compute in practice never actually sums over the vocabulary for the target: it just indexes into one entry.

Why the one-hot target collapses CE to -log p(correct)

Take the cross-entropy sum and plug in the one-hot target. Every y_i is zero except y_c = 1:

H(y, p) = - Σ_i y_i log p_i
        = - (0·log p_1 + ... + 1·log p_c + ... + 0·log p_V)
        = - log p_c

So the entire cross-entropy at one position is just -log p_c: the negative log-probability the model assigned to the correct token. This is identical to the per-token negative log-likelihood from two sections ago — NLL and one-hot cross-entropy are literally the same number. The name you use depends on the lens: likelihood emphasizes the statistical origin, cross-entropy the information-theoretic one.

This collapse is a real computational gift. You never build the one-hot vector, never multiply by zeros, never sum V product terms for the target. You compute the full log p distribution once (via softmax), then gather the single entry at the true token id and negate it. In PyTorch that is F.cross_entropy(logits, targets) or F.nll_loss(log_softmax(logits), targets) — both index rather than dot-product against a one-hot. The target tensor is a vector of integer ids, not a [T, V] matrix, which for a 128k vocabulary is an enormous memory saving.

Entropy, KL divergence, and what CE is really measuring

Cross-entropy has a clean decomposition that explains what the optimizer is actually doing. For any target y and prediction p:

H(y, p) = H(y) + D_KL(y || p)

where H(y) = -Σ_i y_i log y_i is the entropy of the target (its intrinsic uncertainty) and D_KL(y||p) = Σ_i y_i log(y_i / p_i) is the Kullback–Leibler divergence — how far the prediction sits from the target, always ≥ 0 and zero only when p = y. Cross-entropy is thus ‘the irreducible cost of the target plus the penalty for being wrong.’ The optimizer cannot touch H(y) — it depends only on the data, not on θ — so minimizing cross-entropy is exactly minimizing the KL divergence between the target and the model.

With a one-hot target there is a further simplification: a distribution with all its mass on one outcome has zero entropy, H(y) = 0 (since 1·log 1 = 0). So for language-model training the cross-entropy equals the KL divergence outright. Minimizing -log p_c is pulling the model’s whole distribution toward the point mass on the true token — that is the geometric meaning of the loss.

Softmax: turning logits into the probabilities CE consumes

Cross-entropy needs a probability distribution, but the model emits unbounded logits. The softmax is the bridge:

p_i = softmax(z)_i = exp(z_i) / Σ_{j=1}^{V} exp(z_j)

Exponentiating makes every entry positive; dividing by the sum makes them add to 1. The result is a valid distribution that is monotonic in the logits — a larger logit always yields a larger probability — and invariant to adding a constant to every logit, a property we exploit for numerical stability. Softmax and cross-entropy are almost always discussed together because they are two halves of one operation: softmax maps logits to p, cross-entropy scores p against the target. Composed, they form the standard classification head of essentially every transformer.

It matters that this composition is convex in the logits for a fixed target, giving a single well-behaved minimum per example and none of the flat, saturating regions that plague, say, a mean-squared-error-on-sigmoid pairing. But the deeper reason to always fuse them is the gradient. Computed together, softmax-then-cross-entropy has a derivative so simple it looks like a coincidence. It is not a coincidence — it is the entire reason this is the default loss, and it is what the next two sections derive.

The clean gradient: softmax + CE = p - y

The payoff of the softmax–cross-entropy pairing is its gradient with respect to the logits. Let L = -Σ_i y_i log p_i with p = softmax(z). We want ∂L/∂z_k. First, the derivative of a single log-probability through the softmax:

log p_i = z_i - log Σ_j exp(z_j)
∂(log p_i)/∂z_k = δ_{ik} - p_k        (δ_{ik} = 1 if i=k else 0)

Now substitute into the loss and sum over i:

∂L/∂z_k = - Σ_i y_i (δ_{ik} - p_k)
            = - y_k + p_k Σ_i y_i
            = p_k - y_k          (since Σ_i y_i = 1)

So the gradient of the loss with respect to the logits is simply ∇_z L = p - y — the predicted distribution minus the target. For a one-hot target that is p with 1 subtracted at the correct index: every wrong token is pushed down in proportion to how much probability it wrongly received, and the correct token is pushed up by 1 - p_c. No exponentials, no division, no per-term chain rule survive into the final expression. This is the cleanest gradient in deep learning, and it is why softmax + cross-entropy is universal.

Why that pairing is numerically nice

The clean gradient is not just elegant — it is numerically robust, and that robustness comes from fusing the two operations instead of computing them separately. Done naively, softmax exponentiates logits, and a logit of 90 gives exp(90) ≈ 10^39, which overflows float32. Then cross-entropy takes log of the result, and a probability that rounded to 0 gives log 0 = -∞. Two ways to blow up.

The fix is the log-sum-exp identity, applied inside a fused log_softmax. Because softmax is invariant to shifting all logits by a constant, subtract the max logit m = max_j z_j first:

log p_i = z_i - m - log Σ_j exp(z_j - m)

Now every exp argument is ≤ 0, so every term is in (0, 1] — no overflow — and the largest term is exactly 1, so the sum can never underflow to zero and the outer log stays finite. This is why libraries expose F.cross_entropy (which takes raw logits) rather than asking you to call softmax then log then nll yourself: the fused path is both faster and stable. On CPU SLMs, where float32 is common and there is no headroom for silent inf/nan, always feed logits straight to the fused loss — never a hand-rolled softmax.

A worked example, part 1: computing the loss

Concrete numbers make it stick. Take a toy vocabulary of four tokens [a, b, c, d], and suppose the true next token is c (index 2). The model emits logits:

z = [ 2.0,  1.0,  0.1, -1.0 ]      y = [0, 0, 1, 0]  (one-hot at c)

Exponentiate each logit:

exp(z) = [ 7.389, 2.718, 1.105, 0.368 ]
sum    = 7.389 + 2.718 + 1.105 + 0.368 = 11.580

Divide to get the softmax probabilities:

p = [ 0.638, 0.235, 0.0954, 0.0318 ]      (sums to 1.000)

The model’s probability for the correct token c is p_2 = 0.0954 — it actually favored token a. The loss is the negative natural log of that one entry:

L = -log p_2 = -log(0.0954) = 2.349 nats

Sanity check against the bounds: a uniform guess over four tokens would score -log(1/4) = log 4 = 1.386 nats. Our loss of 2.349 is worse than uniform — correct, because the model put most of its mass on the wrong token. A confident correct guess (p_2 → 1) would drive the loss toward 0. The single scalar 2.349 is exactly what backprop will differentiate.

A worked example, part 2: computing the gradient

Now the gradient, which by the derivation is just p - y. Subtract the one-hot target from the probability vector — only the correct index changes by a full unit:

∇_z L = p - y
       = [0.638, 0.235, 0.0954, 0.0318] - [0, 0, 1, 0]
       = [ +0.638, +0.235, -0.905, +0.0318 ]

Read the signs. The correct token c has a negative gradient (-0.905): gradient descent will increase its logit next step. Every wrong token has a positive gradient: their logits get pushed down, and the push is largest for the token the model wrongly liked most (a, +0.638) and smallest for the one it already disfavored (d, +0.0318). The update is perfectly proportional to the error.

Notice the gradient sums to zero: 0.638 + 0.235 - 0.905 + 0.0318 ≈ 0. That is guaranteed, because Σ_k (p_k - y_k) = Σ p_k - Σ y_k = 1 - 1 = 0. It reflects the softmax constraint that probabilities are coupled — raising one must lower others — and it is a cheap invariant to assert in a custom backward pass when you are debugging a hand-written kernel on CPU.

Perplexity: cross-entropy you can actually feel

Cross-entropy is measured in nats, which is hard to have intuition about. Perplexity is the same information exponentiated back into ‘effective number of choices’:

PPL = exp(CE)        (natural log)    or    PPL = 2^CE  (log base 2)

A perplexity of K means the model is, on average, as uncertain as if it were choosing uniformly among K equally likely tokens. For our single worked example, PPL = exp(2.349) = 10.48 — the model was as confused as a uniform guess over ~10.5 tokens, which also equals 1 / p_c = 1 / 0.0954, a tidy identity for a single one-hot prediction.

Because exp is monotonic, minimizing cross-entropy and minimizing perplexity are the same optimization — perplexity is purely a reporting transform. It is the standard headline metric for language models because it is comparable and interpretable: a perplexity of 20 versus 40 immediately says one model halves the effective branching factor. One caveat that trips people up: perplexity depends on the tokenizer. A model over characters and a model over BPE subwords are not comparable by raw perplexity, because they are predicting different-sized units — only bits-per-byte normalizes across tokenizations.

Per-token averaging and masking

The derivation gives the loss at one position. A real batch has B sequences of length T, so the reported loss is the mean of the per-token cross-entropies:

L = (1 / N) Σ_{t} [ -log p(x_t | x_<t) ]     N = number of scored tokens

Averaging rather than summing is what makes the loss scale-free: it does not balloon just because the batch or sequence got longer, so a single learning rate keeps working as you change batch size. But which tokens count is a subtle, bug-prone detail. Sequences are padded to a common length, and those padding positions must be excluded — they carry no real target. Libraries handle this with an ignore_index (commonly -100 in PyTorch): tokens with that target id contribute neither to the numerator nor to N.

Two classic mistakes: dividing by the padded length instead of the count of real tokens (which silently shrinks the loss and starves the gradient), and forgetting the causal shift — the target at position t is the input token at t+1, so logits and labels must be offset by one before scoring. Getting either wrong produces a loss that looks plausible but trains the wrong objective, so both deserve a unit test.

Label smoothing: softening the one-hot target

A one-hot target tells the model to drive p_c → 1, which requires the correct logit to run off toward +∞ relative to the rest. That encourages overconfident, poorly calibrated distributions and can hurt generalization. Label smoothing is a small, cheap regularizer that softens the target: instead of putting all mass on the true token, reserve a little for everyone.

y_i = 1 - ε           if i = c        (e.g. ε = 0.1)
y_i = ε / (V - 1)     otherwise

The loss is still cross-entropy -Σ_i y_i log p_i, but now the target is no longer one-hot, so the collapse to a single term no longer applies and the gradient becomes p - y_smooth — the model is pulled toward a slightly-hedged distribution rather than a spike. The practical effect is better calibration and often a small accuracy or BLEU gain, at the cost of a marginally higher raw loss (you are no longer allowed to be perfectly confident). It is a pointer worth knowing rather than a default: many modern LLM recipes skip it, but it appears throughout the machine-translation and vision-transformer literature. Note that a smoothed target has nonzero entropy H(y) > 0, so cross-entropy no longer equals the KL divergence exactly.

CPU-SLM implications: the loss layer in practice

For a small language model trained or fine-tuned on a CPU, the cross-entropy head has a few practical consequences worth internalizing. First, memory: the logits tensor is [B, T, V], and for a large vocabulary that is often the single biggest activation in the whole forward pass — larger than any attention matrix for a small model. Materializing log_softmax over it, then a second full-size buffer for the backward p - y, can dominate the memory budget; fused cross-entropy kernels that never store the full softmax are a real win.

Second, compute: the final [T, d] × [d, V] output projection plus the softmax is a large matmul, often a meaningful fraction of per-token cost when V is big relative to a small d — one reason tied input/output embeddings and modest vocabularies suit CPU SLMs. Third, stability: always use the fused cross_entropy(logits, targets) path, keep the reduction as a masked mean over real tokens, and if you see nan loss, suspect an unmasked pad, a label off by the causal shift, or a hand-rolled softmax that overflowed — not the optimizer. The math is forgiving; the plumbing is where the bugs live.

Putting the whole chain together

Step back and the derivation is one straight line. Maximum likelihood says maximize ∏_t p_θ(x_t | x_<t). Logarithm turns the product into a sum. Negation turns maximize into minimize, giving the negative log-likelihood -Σ_t log p_θ(x_t | x_<t). Recognizing each term as a cross-entropy against a one-hot target rewrites it as -log p_c per position. Softmax supplies p from the logits, and their fused gradient is the clean p - y.

Every piece has a job: likelihood gives the goal, the log gives numerical sanity, the one-hot target gives the collapse to a single term, softmax gives a valid distribution, log-sum-exp gives stability, and the p - y gradient gives an update proportional to the error that costs almost nothing to compute. Perplexity re-expresses the same loss as an effective vocabulary size for reporting; per-token averaging and masking make it a well-scaled objective across batches; label smoothing is the optional dial on target sharpness. Understand this chain once and the rest of training — why the loss curve looks the way it does, why a broken mask corrupts it, why perplexity is the metric — stops being folklore and becomes arithmetic you can rederive from the one principle: make the observed text probable.