Softmax Classification

Linear regression predicts a number. But most questions aren't how much — they're which category: cat or dog, spam or inbox, one of ten garments. The plumbing barely changes. We just need the output to be a probability distribution over classes, and a loss that rewards putting mass on the truth.

The whole operation, in two formulas — squash logits into probabilities, then score them by the surprise of the correct class:

That log is doing real work. It's gentle when you're right and unbounded when you're confidently wrong.

One prediction is easy to score. The real job is turning every prediction right — moving the weights until the boundaries between classes land where the data says they should.

That live run hid something dangerous. Exponentiating raw logits is a fast path to numerical disaster — and the fix is the reason frameworks ask for logits, not probabilities.

The lab

Two claims to poke at, yours to break: that the classifier really reads pixels — so corrupting them fools it — and that softmax has a temperature you can turn.

How good is "good"?

Accuracy on the training set is worthless — a model can memorize it. What we report is accuracy on a held-out test set, and that number is itself an estimate with error bars. The book works out how many examples you need: to be 95% sure your test error is within ±0.01 of the truth, about 10,000 — which is exactly why so many benchmarks ship a 10,000-image test split.

When the world shifts

Even a perfectly trained, well-tested model can fail the moment it meets the real world, because the world doesn't hold still. The book names three ways the distribution can drift:

The real code

Every act above carries a </> chip with its section's exact PyTorch. The heart of it — softmax and cross-entropy from scratch — is just a few lines:

def softmax(X):
    X_exp = torch.exp(X)
    partition = X_exp.sum(1, keepdims=True)
    return X_exp / partition            # broadcasting
 
def cross_entropy(y_hat, y):
    return -torch.log(y_hat[list(range(len(y_hat))), y]).mean()

# Fashion-MNIST: 60,000 train / 10,000 test, 28×28, ten classes
data = d2l.FashionMNIST(batch_size=256)
 
# from scratch: a 784×10 weight matrix + softmax forward pass
class SoftmaxRegressionScratch(d2l.Classifier):
    def __init__(self, num_inputs, num_outputs, lr, sigma=0.01):
        self.W = torch.normal(0, sigma, size=(num_inputs, num_outputs),
                              requires_grad=True)
        self.b = torch.zeros(num_outputs, requires_grad=True)
 
    def forward(self, X):
        X = X.reshape((-1, self.W.shape[0]))
        return softmax(torch.matmul(X, self.W) + self.b)
 
# accuracy: fraction of argmax predictions that match the label
def accuracy(self, Y_hat, Y, averaged=True):
    preds = Y_hat.argmax(axis=1).type(Y.dtype)
    return (preds == Y.reshape(-1)).type(torch.float32).mean()
 
# concise: let the framework fold softmax into a numerically stable loss
class SoftmaxRegression(d2l.Classifier):
    def __init__(self, num_outputs, lr):
        self.net = nn.Sequential(nn.Flatten(), nn.LazyLinear(num_outputs))
 
def loss(self, Y_hat, Y, averaged=True):
    return F.cross_entropy(Y_hat, Y.reshape((-1,)))   # Y_hat are logits
 
model = SoftmaxRegression(num_outputs=10, lr=0.1)
trainer = d2l.Trainer(max_epochs=10)
trainer.fit(model, data)

Three things to remember

1. Softmax turns scores into a distribution; cross-entropy scores the distribution. Exponentiate to make everything positive, normalize to make it sum to one, then measure the loss as the surprise of the true class — minus its log-probability.
2. The gradient is prediction minus truth. Softmax and cross-entropy are built for each other: their combined gradient is just ŷ − y, the same clean signal as linear regression, which is what makes classifiers so easy to train.
3. Trust the test set, but only barely. A held-out score is an estimate that sharpens like 1/√n, so ±0.01 needs ~10,000 examples — and it stops being a true test the moment you reuse it to choose a model.

Adapted from Chapter 4 of Dive into Deep Learning by Zhang, Lipton, Li, and Smola (CC BY-SA 4.0). The figures run this site's verified TypeScript port; all computation happens live in your browser.