Cross-Entropy and the Softmax Gradient: One Loss for Every Classifier

The router so far answers one question: is this about serving, yes or no? The Notebook Assistant actually needs to send each question to one of four parts of the blog: serving, architecture, agents, or foundations. That means one probability per route, all four summing to one, and a loss that scores the whole distribution.

The answers are the softmax and cross-entropy. They’re worth a post of their own because they are the output layer and the loss of every language model. A language model is this post’s four-way classifier with 128,256 classes instead of four, applied once per token. Everything here carries over unchanged:

• where the loss comes from (information theory, not convenience);
• why its gradient is again “prediction minus target”, the same cancellation Learning Is Optimization found for the sigmoid;
• how to compute it without overflowing;
• how label smoothing changes it;
• what it costs at language-model scale.

The four-way router is the running test throughout. Its mistakes, read from a 4×4 confusion matrix, show what a bag of words can and can’t tell apart.

Table of contents

From one sigmoid to four routes

Learning Is Optimization’s router produced one logit $z$ and turned it into a probability with the sigmoid. For $K$ classes, give the model one logit per class, $z \in \mathbb{R}^K$, and turn them into a distribution with the softmax. (In this post $K$ is the number of classes and $H(\cdot)$ is entropy, both local; they aren’t the key/value heads and head dimension of the notation in AI Engineering, from the Inside.)

The Residual Stream works through the softmax with numbers. Two of its properties matter here.

• Only differences between logits matter. Adding the same constant to every logit changes nothing.
• The scale sets the sharpness. Dividing every logit by a temperature $\tau$ flattens the distribution for $\tau > 1$ and sharpens it for $\tau < 1$. Measuring a Classifier Honestly used exactly this to calibrate the router.

The sigmoid isn’t a separate idea. It’s the softmax with two classes. Fix the second logit at 0, since only the difference matters:

So the binary router was a two-class softmax all along, with the “not serving” logit pinned at zero.

Backpropagation, by Hand found that a hidden layer only ties the linear router at this data size, so the four-way router stays linear, on Features and Overfitting’s bag of words. $W$ grows from one column to four, one per route, and the model computes $z = x^\top W + b$ with $b \in \mathbb{R}^4$. Each route gets its own weights for every word, and the softmax makes the routes compete for probability.

Why not keep four separate sigmoids, one “is it this route?” classifier each (one-vs-rest)? one_vs_rest.py trains both on 20 random splits. The softmax model is the four-way training loop from router4.py, shown in Four routes below; the one-vs-rest model is four sigmoids trained side by side:

def one_vs_rest(X, y, lam=0.003, steps=4000, eta=0.5):
    """Four independent sigmoid classifiers, one per route."""
    W, b = np.zeros((X.shape[1], 4)), np.zeros(4)
    Y = np.eye(4)[y]
    for _ in range(steps):
        P = 1 / (1 + np.exp(-(X @ W + b)))
        G = (P - Y) / len(y)                            # the same p - y, per column
        W -= eta * (X.T @ G + lam * W)
        b -= eta * G.sum(axis=0)
    return W, b

Full script: one_vs_rest.py

val accuracy: softmax 0.621, one-vs-rest 0.624 (difference +0.003 ± 0.016)
one-vs-rest scores: sum is off from 1 by 0.19 on average; no route above 0.5 on 53% of questions (softmax: 33%)
val CE: softmax 0.891, one-vs-rest rescaled to sum to 1 0.885 (difference -0.006 ± 0.011)

The two tie on accuracy. The difference is what the scores mean. Four independent sigmoids can all be low or all be high. Here their sum is off from 1 by 0.19 on average, and on 53% of questions no route clears 0.5, against 33% for the softmax. The scores can be rescaled to sum to 1, and here that scores the same as the softmax: validation cross-entropy 0.885 against 0.891, a tie. But then the distribution in use is a fix applied after training, not the thing training optimized. When exactly one answer is right, as with a route or the next token, the softmax trains the distribution that is actually read. (Per-class sigmoids aren’t wrong in general. They’re the right model when several labels can be true at once. SigLIP, a large-scale image–text contrastive model, also uses them on purpose: each image–text pair is scored on its own, with no softmax over the batch.)

How surprised should the model be?

The binary router was trained on log-loss, $-\log p$ of the true label, motivated in Learning Is Optimization as the Bernoulli likelihood. For $K$ classes the same reasoning gives $-\log p_y$, the negative log-probability of the true class. It has a second derivation, from information theory, that explains what the number means and why it’s the right loss for predicting text.

Surprise. Call $-\log p$ the surprise of an event with probability $p$. A certain event ($p = 1$) has surprise 0. A rare one has large surprise. The log is the one continuous choice, up to its base, that makes surprise add up. Two independent events have probability $p_1 p_2$, and $-\log(p_1 p_2) = -\log p_1 - \log p_2$. That’s also why a dataset’s likelihood, a product over examples, becomes a loss that is a sum over examples. With natural logs the unit is nats; with $\log_2$ it’s bits, and $1$ nat $= 1.443$ bits.

Entropy is the expected surprise when outcomes follow a distribution $p$:

It’s the unavoidable uncertainty in the outcome itself. A fair four-way choice has $H = \log 4 = 1.386$ nats; a certain one has $H = 0$. Shannon showed it’s also a limit on compression: no code for outcomes drawn from $p$ can average fewer than $H(p)$ nats per outcome, and codes over long blocks of outcomes get arbitrarily close.

Cross-entropy is the expected surprise when outcomes follow $p$ but the predictions come from a model $q$:

The second form is the key. The KL divergence $\operatorname{KL}(p \| q) = \sum_k p_k \log(p_k / q_k)$ is never negative, and it’s zero only when $q = p$. So cross-entropy is the true uncertainty plus a penalty for being wrong about it. Minimizing it over $q$ can only lower the KL term, because $H(p)$ doesn’t depend on the model. The best possible model predicts the true distribution, and then it pays exactly the entropy, no more.

With a single labeled example, $p$ is one-hot on the true class $y$, and the sum collapses to $-\log q_y$. That’s the loss we train on, averaged over examples. In this section $p$ is the truth and $q$ the model. From the next section on, as in the rest of the series, $p$ is the model’s softmax output and $y$ is the target.

Four models predicting one question whose true route distribution is $(0.7, 0.1, 0.1, 0.1)$ show what the decomposition measures:

import numpy as np

def H(p):                  # entropy, in nats
    p = np.asarray(p, float)
    return -(p[p > 0] * np.log(p[p > 0])).sum()

def CE(p, q):              # cross-entropy of model q on true distribution p
    p, q = np.asarray(p, float), np.asarray(q, float)
    return -(p[p > 0] * np.log(q[p > 0])).sum()

truth = [0.7, 0.1, 0.1, 0.1]     # what the answer really is, for some question
for name, q in {"matches the truth": truth,
                "uniform guess":     [0.25] * 4,
                "overconfident":     [0.97, 0.01, 0.01, 0.01],
                "confidently wrong": [0.01, 0.97, 0.01, 0.01]}.items():
    ce = CE(truth, q)
    print(f"{name:18s} CE {ce:.3f} = H {H(truth):.3f} + KL {ce - H(truth):.3f}   "
          f"perplexity {np.exp(ce):.2f}")

print(f"\nuniform over 4 routes: ln 4 = {np.log(4):.3f} nats; "
      f"over a 128,256-token vocabulary: {np.log(128256):.2f} nats")
print(f"val CE 0.874 -> perplexity {np.exp(0.874):.2f}")

# Label smoothing: the best a model can do is to predict the smoothed target,
# which caps the gap between the right logit and the others.
for eps in (0.1, 0.01):
    top, rest = 1 - eps + eps / 4, eps / 4
    print(f"smoothing {eps}: target ({top:.4f}, {rest:.4f} x3), logit gap ln({top:.4f}/{rest:.4f}) "
          f"= {np.log(top / rest):.2f}, loss floor {H([top, rest, rest, rest]):.3f}")

Full script: info.py

matches the truth  CE 0.940 = H 0.940 + KL 0.000   perplexity 2.56
uniform guess      CE 1.386 = H 0.940 + KL 0.446   perplexity 4.00
overconfident      CE 1.403 = H 0.940 + KL 0.462   perplexity 4.07
confidently wrong  CE 4.148 = H 0.940 + KL 3.207   perplexity 63.29

uniform over 4 routes: ln 4 = 1.386 nats; over a 128,256-token vocabulary: 11.76 nats
val CE 0.874 -> perplexity 2.40
smoothing 0.1: target (0.9250, 0.0250 x3), logit gap ln(0.9250/0.0250) = 3.61, loss floor 0.349
smoothing 0.01: target (0.9925, 0.0025 x3), logit gap ln(0.9925/0.0025) = 5.98, loss floor 0.052

Read the first four lines.

• The truth itself pays 0.940 nats, the entropy. Nothing can pay less.
• The uniform guess pays 1.386, which is $\log 4$ no matter what the truth is. A four-way classifier scoring worse than 1.386 is worse than the uniform guess. With balanced routes like ours, that means worse than knowing nothing. (With unbalanced classes, the no-information baseline is the entropy of the class frequencies, which is lower.)
• Overconfident but right, putting 0.97 on the likely answer, pays more than the uniform guess. It’s right 70% of the time and pays little then, but the 30% of the time it’s wrong it pays $-\log 0.01 = 4.6$. Cross-entropy punishes confidence that the evidence doesn’t support.
• Confidently wrong pays 4.148. The penalty has no ceiling: as the probability on the truth goes to 0, the loss goes to infinity.

Perplexity is $e^{\text{CE}}$: the loss turned back into “how many equally likely choices would be this uncertain?”. The uniform guess has perplexity 4.00 over four routes. A language model reporting perplexity 10 is, on average, as unsure as picking among 10 equally likely tokens. The four-way router below reaches a test cross-entropy of 0.848, perplexity 2.34: about as unsure as a choice between two or three routes.

The gradient is still prediction minus target

Learning Is Optimization found that the sigmoid and log-loss combine so that the gradient with respect to the logit is $p - y$. The softmax and cross-entropy do the same, and the derivation takes three lines.

The softmax’s derivative. Differentiate $p_i = e^{z_i} / \sum_j e^{z_j}$ with respect to logit $z_j$. For $i = j$ the quotient rule gives $p_i(1 - p_i)$; for $i \neq j$ it gives $-p_i p_j$. Both cases are one formula, with $\delta_{ij} = 1$ if $i = j$ and 0 otherwise:

Raising one logit increases its own probability and takes mass from every other class, in proportion to how much each had.

The chain rule. Write the loss for a target distribution $y$ (one-hot or not) as $\mathcal{L} = -\sum_i y_i \log p_i$, so $\partial\mathcal{L} / \partial p_i = -y_i / p_i$. Then

because the target sums to 1. The $p_i$ from the loss cancels the $p_i$ in the softmax’s derivative, and all that’s left is prediction minus target, one component per class. The derivation never assumed $y$ was one-hot, so it holds for any target distribution, including the smoothed targets below.

Three things follow.

• Every wrong class is pushed down by exactly its own probability. A class the model gave 0.3 gets gradient $+0.3$; the true class gets $p_y - 1$. The loss spends its effort where the model’s mass is misplaced.
• The gradient sums to zero across classes, since both $p$ and $y$ sum to 1. That matches the shift invariance: moving every logit up together changes nothing, so the gradient has no component in that direction.
• It never saturates when the model is wrong. If $p_y \to 0$, the true class’s gradient approaches $-1$, its largest value. Learning Is Optimization showed MSE on a sigmoid stalling exactly there, because its gradient carried an extra $p(1-p)$ factor. The log removes that factor.

For a batch of $n$ examples, stack the gradients as rows: $G = (P - Y)/n$ with $P$ the softmax outputs and $Y$ the targets. Then the rules from Backpropagation, by Hand give the weight gradient $X^\top G$, the bias gradient as the column sums of $G$, and the gradient passed back to whatever produced $X$ as $G W^\top$. In a language model, that last one is where backpropagation into the network begins.

As in Backpropagation, by Hand, check the formula against finite differences before trusting it:

import numpy as np

def softmax(Z):
    """Row-wise softmax, shifted by the row max so exp never overflows."""
    Z = Z - Z.max(axis=-1, keepdims=True)
    E = np.exp(Z)
    return E / E.sum(axis=-1, keepdims=True)

def cross_entropy(Z, y):
    """Mean of -log p[y], computed from logits with log-sum-exp."""
    Z = Z - Z.max(axis=-1, keepdims=True)
    log_p = Z - np.log(np.exp(Z).sum(axis=-1, keepdims=True))
    return -log_p[np.arange(len(y)), y].mean()

if __name__ == "__main__":
    # The gradient of cross-entropy w.r.t. the logits is p - onehot(y).
    rng = np.random.default_rng(0)
    z, target = rng.normal(size=(1, 4)), np.array([2])
    analytic = softmax(z)[0] - np.eye(4)[2]
    numeric = np.array([(cross_entropy(z + e, target) - cross_entropy(z - e, target)) / 2e-6
                        for e in np.eye(4)[:, None, :] * 1e-6])
    print(f"p - y:  {np.round(analytic, 4)}")
    print(f"largest gap to finite differences: {np.abs(numeric - analytic).max():.1e}")

    # Why the shift matters.
    z = np.array([[1000.0, 999.0, 0.0, -5.0]])
    with np.errstate(over="ignore", invalid="ignore"):
        naive = np.exp(z) / np.exp(z).sum()
    print(f"naive softmax: {naive[0]}   shifted: {np.round(softmax(z)[0], 4)}")
    show = lambda v: "[" + " ".join(f"{x:.2f}" for x in v[0]) + "]"
    with np.errstate(divide="ignore"):
        print(f"log(softmax(z)): {show(np.log(softmax(z)))}   "
              f"log-sum-exp: {show(z - z.max() - np.log(np.exp(z - z.max()).sum()))}")

Full script: softmax_basics.py

p - y:  [ 0.226   0.1746 -0.6219  0.2213]
largest gap to finite differences: 8.5e-11
naive softmax: [nan nan  0.  0.]   shifted: [0.7311 0.2689 0.     0.    ]
log(softmax(z)): [-0.31 -1.31 -inf -inf]   log-sum-exp: [-0.31 -1.31 -1000.31 -1005.31]

The later scripts import softmax and cross_entropy from softmax_basics.py. The analytic gradient matches the numerical one to $10^{-10}$, and its four entries sum to zero. The true class, index 2, is the only negative entry.

Computing it without overflow

The last two lines of that output are the reason for the two max shifts in softmax_basics.py. With logits $[1000, 999, 0, -5]$, np.exp(1000) overflows to infinity, and $\infty / \infty$ is nan. Subtracting the largest logit first gives $[0, -1, -1000, -1005]$: the same softmax, and nothing larger than $e^0 = 1$.

The loss needs one more step. Computing softmax and then log can still fail in the other direction. The shifted softmax gives the last two routes probabilities around $e^{-1000}$, which underflow to exactly 0, and $\log 0 = -\infty$: the last line’s log(softmax(z)). If one of those were the true class, the loss would be infinite and its gradient nan. So compute the log-probabilities directly from the logits with the log-sum-exp identity, where $m = \max_k z_k$:

After the shift, the largest term in the sum is $e^0 = 1$, so the sum is between 1 and $K$. Its log can neither overflow nor be $\log 0$, and the output’s last line shows the result: $-1000.31$ and $-1005.31$, large but finite. That’s what cross_entropy in softmax_basics.py does, and it’s why deep-learning libraries take logits, not probabilities, as input to their cross-entropy. PyTorch’s F.cross_entropy(logits, target) applies the softmax, the log-sum-exp and the loss in one numerically safe step. Its target can be class indices or, since version 1.10, full class probabilities, which fits a gradient that works for any target distribution, and it has a label_smoothing argument for the next section. Passing it probabilities from an earlier softmax is a common bug: it applies log_softmax to them, a second softmax.

Four routes

The four-way router uses Features and Overfitting’s bag of words, with each question labeled by its route index. Training is Learning Is Optimization’s gradient descent with the new gradient:

import numpy as np
from router_data import featurize, ROUTES
from softmax_basics import softmax

def four_way(train, *others):
    """Bag of words from "Features and Overfitting", with the route index (0-3) as the label."""
    sets, vocab = featurize(train, *others)
    return [(X, np.array([ROUTES.index(r["route"]) for r in rs]))
            for (X, _), rs in zip(sets, (train, *others))], vocab

def train(X, y, lam=0.0, smooth=0.0, steps=4000, eta=0.5, K=4):
    W, b = np.zeros((X.shape[1], K)), np.zeros(K)
    Y = np.eye(K)[y] * (1 - smooth) + smooth / K        # label smoothing
    for _ in range(steps):
        G = (softmax(X @ W + b) - Y) / len(y)           # p - y, per example
        W -= eta * (X.T @ G + lam * W)
        b -= eta * G.sum(axis=0)
    return W, b

def top_label_ece(P, y, bins=5):
    conf, right = P.max(axis=1), P.argmax(axis=1) == y
    which = np.clip((conf * bins).astype(int), 0, bins - 1)
    return sum((which == k).mean() * abs(conf[which == k].mean() - right[which == k].mean())
               for k in range(bins) if (which == k).any())

Full script: router4.py

(router_data.py is from Features and Overfitting. top_label_ece is Measuring a Classifier Honestly’s calibration error applied to the model’s top prediction: how far “80% sure” is from being right 80% of the time.)

With four classes, the label smoothing in train needs its own explanation. Label smoothing replaces the one-hot target with a mixture: $(1 - \varepsilon)$ on the true class plus $\varepsilon$ spread evenly over all $K$. With $\varepsilon = 0.1$ and $K = 4$ the target becomes $(0.925, 0.025, 0.025, 0.025)$. Since the gradient is $p - y$ for any target, the model is now pulled toward 0.925, not 1.0, and the effect shows up in two places (from the last two lines of info.py’s output):

• A loss floor. The best achievable loss is the smoothed target’s entropy, 0.349 nats, not 0.
• Bounded logits. Matching the target means the true class’s logit sits exactly $\log(0.925/0.025) = 3.61$ above the others. Without smoothing, the only way to reach the target is an infinite gap, so on data it can separate, the model keeps growing its weights. That’s the failure Learning Is Optimization predicted and Features and Overfitting found. A smaller $\varepsilon$ allows a bigger gap and a lower floor, 5.98 and 0.052 at $\varepsilon = 0.01$, and as $\varepsilon \to 0$ the gap is unbounded again.

Smoothing is a regularizer on the output rather than on the weights. The smoothed loss splits into two cross-entropies, $(1 - \varepsilon)$ times the usual one plus $\varepsilon$ times the cross-entropy against the uniform distribution, and that second term is a penalty on confident outputs.

The experiment compares four models on Features and Overfitting’s 20 random splits, 30 training, 10 validation and 10 test questions per route. It uses Features and Overfitting’s method throughout: differences are paired split by split, with the Nadeau–Bengio corrected standard error, and the choice is made on validation cross-entropy:

CONFIGS = {"L2 0.003": (0.003, 0.0), "no penalty": (0.0, 0.0),
           "smoothing 0.1": (0.0, 0.1), "smoothing 0.1 + L2 0.003": (0.003, 0.1)}
val, test = {c: [] for c in CONFIGS}, {c: [] for c in CONFIGS}
for seed in range(20):
    ((Xtr, ytr), (Xva, yva), (Xte, yte)), _ = four_way(*split(seed))
    for name, (lam, smooth) in CONFIGS.items():
        W, b = train(Xtr, ytr, lam, smooth)
        P = softmax(Xva @ W + b)
        val[name].append((cross_entropy(Xva @ W + b, yva), (P.argmax(1) == yva).mean(),
                          top_label_ece(P, yva)))
        test[name].append((softmax(Xte @ W + b).argmax(1), yte, cross_entropy(Xte @ W + b, yte)))

The script then prints each model’s mean against L2’s, paired split by split, and for the model chosen on validation cross-entropy, the test confusion matrix pooled over the 20 splits.

Full script: four_routes.py

                           val CE           vs L2  val acc           vs L2    ECE           vs L2
L2 0.003                    0.891                    0.621                  0.096                
no penalty                  0.976  +0.085 ± 0.054    0.633  +0.011 ± 0.021  0.154  +0.059 ± 0.033
smoothing 0.1               0.874  -0.017 ± 0.022    0.655  +0.034 ± 0.029  0.106  +0.010 ± 0.033
smoothing 0.1 + L2 0.003    0.907  +0.016 ± 0.014    0.628  +0.006 ± 0.013  0.125  +0.029 ± 0.031
smoothing + L2 vs smoothing alone: val CE +0.033 ± 0.018

chosen on validation CE: smoothing 0.1. test CE 0.848 (perplexity 2.34)
test confusion, pooled over the 20 splits (rows = true route):
                    serving architecture       agents  foundations
serving                 120           29           22           29
architecture             19          113           12           56
agents                   10           10          153           27
foundations              17           19           20          144
serving        precision 0.723  recall 0.600  F1 0.656
architecture   precision 0.661  recall 0.565  F1 0.609
agents         precision 0.739  recall 0.765  F1 0.752
foundations    precision 0.562  recall 0.720  F1 0.632
test accuracy 0.6625; mean recall 0.6625; macro-F1 0.6621

The top half is the comparison.

• No penalty against L2: the metrics disagree. L2’s validation cross-entropy is better by 0.085 ± 0.054, and its calibration error by 0.059 ± 0.033, both suggestive rather than decisive. On accuracy the unpenalized model is 0.011 ahead, well inside its error. The two metrics measure different things. Accuracy only checks the top route. Cross-entropy also charges for how confident each prediction was, and without a penalty the probabilities run hotter, including on the questions it gets wrong.
• Smoothing and L2 tie. Smoothing’s validation cross-entropy is lower by 0.017 ± 0.022 and its accuracy is higher by 0.034 ± 0.029. Both are within about one standard error, so neither is a real difference at this data size.
• Both together are too much. Combined, they’re 0.033 ± 0.018 worse than smoothing alone and tie L2.

Features and Overfitting warned that a fixed step count can quietly act as early stopping, and that an unpenalized model’s weights keep growing. Both are worth checking, along with what each model does on its own training data:

rows = {}
for seed in range(20):
    ((Xtr, ytr), (Xva, yva), _), _ = four_way(*split(seed))
    for name, (lam, smooth) in CONFIGS.items():
        for steps in (4000, 16000):
            W, b = train(Xtr, ytr, lam, smooth, steps=steps)
            P = softmax(Xva @ W + b)
            rows.setdefault((name, steps), []).append((
                cross_entropy(Xva @ W + b, yva), np.linalg.norm(W),
                cross_entropy(Xtr @ W + b, ytr),              # hard labels, no penalty
                P.max(1).mean(), (P.argmax(1) == yva).mean()))

Full script: longer_training.py

                           steps  val CE           vs L2    |W|  train CE  mean top p  val acc
L2 0.003                   4,000   0.891                   10.6     0.131       0.622    0.621
L2 0.003                  16,000   0.891                   10.6     0.131       0.622    0.621
no penalty                 4,000   0.976  +0.085 ± 0.054   21.7     0.008       0.776    0.633
no penalty                16,000   1.068  +0.177 ± 0.078   27.3     0.002       0.815    0.630
smoothing 0.1              4,000   0.874  -0.017 ± 0.022   14.1     0.080       0.616    0.655
smoothing 0.1             16,000   0.876  -0.015 ± 0.025   14.7     0.078       0.613    0.661
smoothing 0.1 + L2 0.003   4,000   0.907  +0.016 ± 0.014    9.0     0.200       0.563    0.628
smoothing 0.1 + L2 0.003  16,000   0.907  +0.016 ± 0.014    9.0     0.200       0.563    0.628

• The unpenalized model was being early-stopped. Trained four times longer, its weight norm grows from 21.7 to 27.3, its validation cross-entropy worsens to 1.068, and L2’s lead doubles to 0.177 ± 0.078. The other three barely move, so for them 4,000 steps is enough.
• Smoothing bounds the weights, as predicted. Without any weight penalty, its weight norm levels off (14.1, then 14.7), while the unpenalized model’s keeps climbing.
• The combination underfits. Its training cross-entropy is 0.200, against 0.080 for smoothing alone, and its average top probability, 0.563, is below its accuracy of 0.628: it’s less sure than it should be. Each method restrains confidence, and together they restrain it too much.

The rule from Features and Overfitting is to choose on validation and not relitigate, so the router uses smoothing 0.1, the lowest validation cross-entropy. It’s chosen on validation, but by a margin smaller than its own error. The test numbers below describe this model, not a claim that smoothing beats L2.

Reading a 4×4 confusion matrix

A two-class router has four outcomes; a four-class router has sixteen. Measuring a Classifier Honestly’s binary definitions carry over one route at a time. Treat route $k$ as “positive” and the other three as “negative”:

• precision for $k$ is its diagonal entry divided by its column sum, of the questions sent to $k$, the fraction that belonged there;
• recall for $k$ is the diagonal entry over its row sum, of the questions that belonged to $k$, the fraction that got there.

The bottom half of the output is the chosen router on the pooled test splits, 200 questions per route.

• Accuracy is 0.662, against 0.25 for guessing, and test cross-entropy is 0.848. With the same number of test questions in every route, accuracy equals the mean of the four recalls exactly: 0.6625 both. Macro-F1, the unweighted mean of the four routes’ F1 scores, also mixes in precision, and lands at 0.6621, close by near-coincidence. With unbalanced routes the three can differ a lot, and macro-F1 weights a rare route as much as a common one.
• Agents is the easiest route, with recall 0.765 and precision 0.739. Its vocabulary, like agent, tool, prompt and context, rarely appears in the other routes’ questions.
• The worst cell is off the diagonal. Of 200 architecture questions, 56 go to foundations, more than to serving and agents together (31). Foundations has the lowest precision, 0.562, because it absorbs about twice as many mistakes as any other route: 112 questions wrongly sent there, against 46 to 58 for the others.

The same matrix as fractions of each route’s 200 test questions, true route down the side and predicted route along the top. The diagonal is what the router gets right. The largest cell off it, 0.28, is architecture sent to foundations:

Why foundations? A look at one split’s learned weights and its mistakes:

import numpy as np
from router_data import split, tokens, ROUTES
from softmax_basics import softmax
from router4 import four_way, train

tr, va, te = split(0)
((Xtr, ytr), (Xte, yte)), vocab = four_way(tr, te)
W, b = train(Xtr, ytr, smooth=0.1)
for k, route in enumerate(ROUTES):
    top = np.argsort(-W[:, k])[:6]
    print(f"{route:13s} " + ", ".join(f"{vocab[j]} {W[j, k]:+.2f}" for j in top))
print(f"\nbias b = {np.round(b, 2)}, softmax(b) = {np.round(softmax(b), 2)}")
P = softmax(Xte @ W + b)
for i in np.where((yte == 1) & (P.argmax(1) == 3))[0]:
    words = tokens(te[i]["question"])
    print(f"\narchitecture -> foundations: {te[i]['question']}")
    print("  never seen in training: " + ", ".join(t for t in words if t not in vocab))
    for t in dict.fromkeys(t for t in words if t in vocab):
        print(f"  {t:5s} weights {np.round(W[vocab.index(t)], 2)}")
    print(f"  words alone: p = {np.round(softmax(Xte[i] @ W), 2)}   with the bias: p = {np.round(P[i], 2)}")

Full script: read_weights.py

serving       vllm +1.11, batching +1.11, token +0.98, paged +0.96, to +0.95, quantization +0.92
architecture  attention +1.92, unembedding +1.61, residual +1.59, stream +1.29, transformers +1.21, layer +1.17
agents        agent +1.83, prompt +1.63, context +1.47, calling +1.38, rot +1.33, retrieval +1.22
foundations   backpropagation +1.74, confusion +1.29, learning +1.23, gradients +1.23, sigmoid +1.12, optimizer +1.09

bias b = [ 0.13 -0.17 -0.17  0.21], softmax(b) = [0.28 0.21 0.21 0.3 ]

architecture -> foundations: How does rotary position embedding encode distance?
  never seen in training: rotary, position, embedding, encode, distance
  how   weights [-0.7   0.19  0.27  0.24]
  does  weights [ 0.43 -0.2  -0.01 -0.23]
  words alone: p = [0.19 0.24 0.32 0.25]   with the bias: p = [0.22 0.21 0.27 0.31]

architecture -> foundations: How is the embedding table trained?
  never seen in training: embedding, table, trained
  how   weights [-0.7   0.19  0.27  0.24]
  is    weights [ 0.04 -0.24 -0.15  0.34]
  the   weights [ 0.02  0.57 -0.43 -0.16]
  words alone: p = [0.12 0.38 0.16 0.34]   with the bias: p = [0.13 0.31 0.14 0.41]

The weights look right: each route’s strongest words are its own topic. The two mistakes show what goes wrong. In both, every content word, rotary, position, embedding, table, trained, never appeared in the 120 training questions, so the bag of words drops them. What’s left is a few function words, how, does, is, the, and the bias. Their evidence is weak, and the bias decides. In this split the bias itself leans toward foundations, 0.30 against 0.21 for architecture. On the words alone, the second question would have gone to architecture, correctly.

The routes are balanced, 30 training questions each, so the bias isn’t a learned prior about how common each route is. It absorbs whatever the word weights leave over, and which route it favors changes from split to split. Across all 20 splits:

unsure, sent_to, all_to, unseen = [], np.zeros(4), np.zeros(4), {True: [], False: []}
to_found, moved, right = 0, 0, {True: [], False: []}
for seed in range(20):
    tr, va, te = split(seed)
    ((Xtr, ytr), (Xte, yte)), vocab = four_way(tr, te)
    W, b = train(Xtr, ytr, smooth=0.1)
    P = softmax(Xte @ W + b)
    low = P.max(1) < 0.4                                 # the router is barely guessing
    unsure.append(low.mean())
    sent_to += np.bincount(P[low].argmax(1), minlength=4)
    all_to += np.bincount(P.argmax(1), minlength=4)
    for i, r in enumerate(te):
        unseen[bool(low[i])].append(np.mean([t not in vocab for t in tokens(r["question"])]))
        right[bool(low[i])].append(P[i].argmax() == yte[i])
    # Would they still go to foundations if no route had a head start from its bias?
    found = low & (P.argmax(1) == 3)
    to_found += found.sum()
    moved += (softmax(Xte[found] @ W + b.mean()).argmax(1) != 3).sum()
print(f"test questions with top probability < 0.4: {np.mean(unsure):.0%}, "
      f"right {np.mean(right[True]):.0%} of the time (the rest: {np.mean(right[False]):.0%})")
print("where they go:   " + ", ".join(f"{r} {n / sent_to.sum():.0%}" for r, n in zip(ROUTES, sent_to)))
print("all predictions: " + ", ".join(f"{r} {n / all_to.sum():.0%}" for r, n in zip(ROUTES, all_to)))
print(f"share of words never seen in training: unsure {np.mean(unseen[True]):.0%}, "
      f"the rest {np.mean(unseen[False]):.0%}")
print(f"of the {to_found} unsure questions sent to foundations, {moved} go elsewhere with an even bias")
for w in ["what", "is"]:
    print(f"questions containing '{w}': " + ", ".join(
        f"{r} {sum(w in tokens(q['question']) for q in english if q['route'] == r)}" for r in ROUTES))

Full script: unsure.py

test questions with top probability < 0.4: 16%, right 40% of the time (the rest: 71%)
where they go:   serving 18%, architecture 23%, agents 22%, foundations 37%
all predictions: serving 21%, architecture 21%, agents 26%, foundations 32%
share of words never seen in training: unsure 33%, the rest 26%
of the 47 unsure questions sent to foundations, 24 go elsewhere with an even bias
questions containing 'what': serving 11, architecture 19, agents 16, foundations 28
questions containing 'is': serving 13, architecture 17, agents 14, foundations 28

• One test question in six is a near-guess, with top probability below 0.4. Those questions have more words the router has never seen, 33% of their words against 26% for the rest. They’re right 40% of the time, against 71% for the rest.
• Near-guesses lean toward foundations: 37% go there. The router leans that way in general, sending 32% of all questions there, and a little more when it’s unsure.
• The bias explains about half of it. Of the 47 near-guesses sent to foundations, 24 go somewhere else when every route gets the same bias.
• Phrasing explains some of the rest. Foundations questions are more often definitions, “What is gradient descent?”, and 28 of 50 foundations questions contain what and 28 contain is, more than any other route. So what and is carry a little foundations weight.

Neither is a flaw in the loss or the optimizer. When a question’s topic words are unknown, the router has nothing to go on but phrasing and a bias, and it does the best it can with them. A bag of words built from 120 questions can’t know that rotary and embedding are architecture words. Features learned from far more text can, which is what embeddings do (Embeddings and Contrastive Learning).

At language-model scale

Everything above is exactly how a language model is trained. The classes are the tokens of the vocabulary. Llama 3’s has 128K tokens, $V = 128{,}256$ in the released model’s configuration. At every position of every training sequence, the model produces $V$ logits for the next token, and the loss is the cross-entropy against the token that actually came next. Two numbers from this post transfer directly.

• The starting loss. A model whose logits all start near zero predicts uniformly and pays $\log V = 11.76$ nats per token, as info.py printed. With small random logits of standard deviation $\sigma$, it pays about $\log V + \sigma^2/2$: slightly more, never much less. That makes a sanity check for the first step of any training run. A first loss well above $\log V$ means the initial logits are too large. One well below it means the model can see the answer, usually a label leak or targets shifted by one position.
• The gradient. Backpropagation into the network starts from $P - Y$ at every position: the model’s next-token distribution minus a one-hot on the real next token.

The new problem is size. The loss needs a logit for every position and every vocabulary entry:

# Memory for the logits of one training batch, at Llama-3-8B's vocabulary.
V, D = 128_256, 4096
for B, T in [(1, 8192), (8, 8192)]:
    n = B * T                                            # token positions
    logits = n * V * 4                                   # FP32 logits
    print(f"B={B}, T={T}: {n:,} positions x {V:,} = {n * V / 1e9:.2f} B logits; "
          f"FP32 {logits / 2**30:.1f} GiB, and the same again for their gradient")
print(f"the hidden states that produce them, BF16: {8 * 8192 * D * 2 / 2**30:.2f} GiB")
print(f"with activation checkpointing, the saved input of each of 32 layers: "
      f"{32 * 8 * 8192 * D * 2 / 2**30:.0f} GiB")

Full script: logit_memory.py

B=1, T=8192: 8,192 positions x 128,256 = 1.05 B logits; FP32 3.9 GiB, and the same again for their gradient
B=8, T=8192: 65,536 positions x 128,256 = 8.41 B logits; FP32 31.3 GiB, and the same again for their gradient
the hidden states that produce them, BF16: 0.50 GiB
with activation checkpointing, the saved input of each of 32 layers: 16 GiB

Even one sequence of 8,192 tokens has a billion logits, 3.9 GiB in FP32, the precision the loss is usually computed in, plus the same again for their gradient. At a batch of 8 such sequences the logits are 8.4 billion numbers, 31.3 GiB. The hidden states that produced them take 0.50 GiB. A full set of activations saved for backward is larger still, but with activation checkpointing (Backpropagation, by Hand), which keeps only each layer’s input, the whole 32-layer model saves 16 GiB. This one intermediate of the loss is then bigger than everything else the model stores for backward. For smaller models with large vocabularies, the Cut Cross-Entropy authors find the loss can take an order of magnitude more memory than the rest of the model combined.

The $p - y$ form is what makes the fix possible. The gradient at a position needs only that position’s logits and its target, so nothing requires the whole matrix to exist at once. Fused cross-entropy kernels compute logits a chunk at a time, take the loss and $p - y$ for that chunk immediately, push the gradient back through the unembedding, and throw the chunk away.

• Liger Kernel fuses the final matrix multiplication with the loss and processes the input in chunks.
• Cut Cross-Entropy goes further: it computes only the true token’s logit and the log-sum-exp, on the fly, in on-chip memory. Splitting the vocabulary into blocks means the log-sum-exp is built up block by block, with a running maximum (an “online” log-sum-exp). For Gemma 2 (2B), its authors report the loss computation’s memory falling from 24 GB to 1 MB.

The math is the same as the four-route router. Only the order in which it’s computed changes, to fit the hardware. That pattern, the same result computed in an order that avoids storing a large intermediate, comes back in FlashAttention, which uses the same running-maximum trick for attention’s softmax.

Where this leads

• Initialization and Optimizers is about making gradient descent work on deeper networks: initialization, why the output layer’s scale sets that first loss, and the optimizers, momentum and Adam, that took over from plain gradient descent.
• Embeddings and Contrastive Learning uses the softmax and cross-entropy over one positive and many negatives, contrastive learning, to train embeddings. That’s the loss behind retrieval.
• What “Generative” Means trains the first language model on the blog. Its loss is this post’s cross-entropy at every position, and its perplexity is this post’s $e^{\text{CE}}$.
• The same softmax normalizes attention scores, in The Transformer’s Three Attention Patterns, and its temperature is the sampling knob of Decoding.
• KL divergence comes back in Knowledge Distillation, which uses it to compare a student to a teacher, and in RLHF, where it keeps a preference-tuned model close to where it started.

Takeaways

• The softmax turns $K$ logits into a distribution. The sigmoid is the two-class case with one logit fixed at 0. Four separate sigmoids score about as well, but their outputs don’t form a distribution unless rescaled after training.
• Surprise is $-\log p$, the only continuous choice (up to the base) that makes independent events add. Entropy is expected surprise. Cross-entropy is entropy plus the KL divergence from the truth, so minimizing it drives the model toward the true distribution.
• Cross-entropy punishes unjustified confidence: 0.97 on an answer that’s right 70% of the time scores worse than a uniform guess. Perplexity, $e^{\text{CE}}$, reads the loss as a number of equally likely choices.
• The gradient with respect to the logits is $p - y$ for any target distribution. It sums to zero, pushes each wrong class down by its own probability, and doesn’t saturate when the model is wrong.
• Compute the loss from logits with log-sum-exp, never as log(softmax(z)). Libraries take logits for this reason.
• Label smoothing gives a loss floor and a bounded logit gap, a regularizer on the output. Here it tied L2, kept the weights bounded without a weight penalty, and made the probabilities run low. Combining it with L2 underfit.
• Cross-entropy and accuracy can disagree. Read both, and calibration too. Check that a fixed step count isn’t early-stopping the baseline.
• With topic words unknown, a bag-of-words router falls back on phrasing and its bias. Better features, not a better loss, are the fix.
• A language model is this classifier over the vocabulary. Its first loss should be near $\log V$, slightly above, never well below. At scale the logits can outweigh everything else a checkpointed model stores for backward, which fused kernels avoid by never storing them.

References

Start with Cover and Thomas, chapter 2, for entropy, cross-entropy and KL with proofs. Then read Müller, Kornblith and Hinton for what label smoothing does to representations. Read the Cut Cross-Entropy paper for the memory problem at scale.

• Shannon, A Mathematical Theory of Communication (Bell System Technical Journal, 1948): entropy, and the limit it sets on the average length of any code.
• Cover and Thomas, Elements of Information Theory (Wiley, 2nd ed., 2006): entropy, cross-entropy and KL divergence, with the proof that KL is never negative.
• Bridle, Probabilistic Interpretation of Feedforward Classification Network Outputs (1990): the softmax as a network’s output layer, and its name.
• Zhai, Mustafa, Kolesnikov and Beyer, Sigmoid Loss for Language Image Pre-Training (ICCV, 2023): SigLIP, a contrastive model trained with per-pair sigmoids instead of a softmax.
• Szegedy, Vanhoucke, Ioffe, Shlens and Wojna, Rethinking the Inception Architecture for Computer Vision (CVPR, 2016): label smoothing as a regularizer on the output.
• Pereyra, Tucker, Chorowski, Kaiser and Hinton, Regularizing Neural Networks by Penalizing Confident Output Distributions (2017): label smoothing as a penalty on confident outputs.
• Müller, Kornblith and Hinton, When Does Label Smoothing Help? (NeurIPS, 2019): label smoothing improves calibration, and makes a teacher worse for distillation.
• Grattafiori et al. (Meta AI), The Llama 3 Herd of Models (2024): the 128K-token vocabulary.
• Hsu et al., Liger Kernel: Efficient Triton Kernels for LLM Training (2024): fused, chunked linear and cross-entropy kernels.
• Wijmans, Huval, Hertzberg, Koltun and Krähenbühl, Cut Your Losses in Large-Vocabulary Language Models (ICLR, 2025): cross-entropy without materializing the logits; 24 GB to 1 MB for Gemma 2 (2B).
• PyTorch, torch.nn.functional.cross_entropy (documentation): takes logits, with a label_smoothing argument.