Concept25 minBefore this:Logistic Regression Entropy, Cross-Entropy & KL Divergence

Softmax Regression

How does logistic regression handle many classes, and why does every LLM end with it?

The situation

You're sorting fruit into apple, orange or banana. A teammate turns scores into probabilities by dividing each by the total, and then one score goes negative. The interviewer asks:

What is softmax regression, and how is it trained?

Running example: sort 27 fruits into apple, orange or banana from two measurements: length (cm) and "yellowness" (0–1, from a colour sensor). Illustrative data, generated with a fixed seed.

The idea

Think of a vote share. Each class gets a score, and raising e to that score makes it positive and widens the lead. Dividing by the total then gives shares that add up to 100%.

Worked from the scores 2, 1 and −1. Percentages are rounded.

Play

Try to break it. Add 3 to every logit. Then raise the temperature, and then lower it.

From logits to probabilities

Move the scores, watch the distribution

apple logit orange logit banana logit temperature T

cross-entropy if it's really an apple

–

sum of probabilities

–

Things to try. Add the same amount to every logit: nothing changes (the Maths tab shows why). Raise T: the distribution flattens toward uniform. Lower T: it sharpens toward the single biggest logit. That's exactly the "temperature" knob in LLM sampling.

Now say why: why does adding the same amount change nothing, and what does temperature really do?

Maths

In one breath

Softmax regression extends logistic regression to K classes. Each class gets its own weight vector and its own score (a logit); the softmax exponentiates the scores and normalises them, so they're positive and sum to 1. It's trained with cross-entropy, and the gradient is beautifully simple: (predicted probability − one-hot label) × features. With two classes it reduces exactly to the sigmoid. And it's not just classical ML: the last layer of every LLM is softmax regression over its vocabulary.What you'll be able to do after this: derive softmax and its gradient, explain why logits are shift-invariant, implement it without overflow, compare it with one-vs-rest, and connect temperature to LLM sampling.

appleorangebananamaths, numberscorrect / fixerrors, overflow

One score per class, then turn scores into probabilities

Give each class k its own weight vector θ_k. For a fruit with features x, each class gets a score t_k = θ_kᵀx: how much the features look like that class. These raw scores are called logits. They can be any real number, so we need a function that turns K arbitrary numbers into a probability distribution. CS229 points out the two problems with using them directly: θ_kᵀx "is not necessarily within [0, 1]" and "the summation of θ_kᵀx's is not necessarily 1" CS229 §2.3.

Notation: the label is often written one-hot, a vector with a 1 in the true class's slot and 0 elsewhere. A banana is (0, 0, 1).

One principle

Exponentiate the scores, then normalise

CS229 Eqs. 2.9–2.10 CS229 §2.3. "Softmax" because it's a smooth version of "argmax": the biggest score gets most of the probability, but the others keep some.

Only differences matter

Adding a constant to every logit changes nothing

Two consequences. First, the parameters are redundant: add the same vector to every θ_k and predictions don't change, so one class's weights can be fixed at zero without losing anything. Regularization also pins this down. Second, it gives the standard trick for computing softmax without overflow:

Computed live in your browser with 64-bit floats. e¹⁰⁰⁰ overflows to Infinity, and Infinity / Infinity is NaN. e⁻¹⁰⁰⁰ underflows to 0, and 0 / 0 is also NaN. Subtracting the largest logit first makes the biggest exponent e⁰ = 1: same answer, no overflow. Every library does this, and computes log-softmax as t − max − log Σ e^{t − max} (the log-sum-exp trick).

The loss and its gradient

Cross-entropy, and the cleanest gradient in ML

Maximum likelihood again: maximise the probability the model gives to the true class. For one example the loss is just −log of that probability CS229 §2.3:

With φ = softmax of the logits; in vector form, ∂ℓ/∂t = φ − e_y (CS229 Eq. 2.17).

–

Using the trained model from the next section. Bars: predicted probability φ (filled) and the one-hot label (outline). The gradient on each class's logit is φ − label: negative means "raise this score", positive means "lower it". Confident and correct → tiny gradient; confused → big push.

What it learns

Straight-edged regions, one per class

Swipe sideways to see the whole diagram.

27 fruits by length and yellowness. Bananas are long and yellow; apples and oranges differ mainly in colour.
After training (gradient descent on cross-entropy, a light L2 penalty), each point of the plane is assigned to the class with the biggest logit. These are the decision regions.
Region edges are where two logits tie: θ_aᵀx = θ_bᵀx, which is a straight line. So softmax regression is still a linear classifier: K classes, regions bounded by straight lines.
Curved class shapes need engineered features, or hidden layers in front: a neural network is learned features plus this same softmax layer.

Go deeper: two classes = the sigmoid

With K = 2 and logits (t₁, t₂): P(class 1) = e^t₁ / (e^t₁ + e^t₂) = 1 / (1 + e^{−(t₁ − t₂)}) = σ(t₁ − t₂). Binary softmax is exactly logistic regression, with weight vector θ₁ − θ₂ (the redundancy from the section above, in action). In GLM terms, softmax is the canonical response function of the multinomial distribution, just as the sigmoid is for the Bernoulli CS229 §3.

In practice

Two ways to go multi-class

Approach	How	Trade-off
One-vs-rest (OvR)	Train K separate binary logistic regressions ("apple vs not", …), predict the highest	Simple and parallel, but the K probabilities don't sum to 1 and weren't trained together (see Softmax vs one-vs-rest below)
Multinomial / softmax	One model, K weight vectors, a single softmax, a single cross-entropy loss	Coherent probabilities; the standard choice; it's the output layer of neural classifiers and LLMs

scikit-learn's LogisticRegression fits the multinomial (softmax) model for multi-class problems with its default solver, and uses a coefficient matrix W with one row per class; the old liblinear solver "cannot learn a true multinomial (multiclass) model" scikit-learn docs. In PyTorch, it's a nn.Linear layer followed by CrossEntropyLoss, which applies the softmax internally.

Softmax vs one-vs-rest

OvR's probabilities don't add up

Both models trained on the same 27 fruits. OvR runs three independent "this class vs the rest" logistic regressions; their outputs answer three different questions, so they can sum to more or less than 1. Softmax's outputs are one coherent distribution. Both usually pick the same winner; the difference shows up when you use the probabilities (ranking, thresholds, expected cost).

The bridge to LLMs

Every LLM ends with softmax regression

A language model's last step is exactly this lesson. The network turns the context into a hidden vector (4,096 numbers in Llama 3.1 8B), and a final linear layer produces one logit per token in the vocabulary: 128,256 of them. Softmax turns those into next-token probabilities, and training minimises cross-entropy against the actual next token Llama 3.1 8B config. That one output matrix alone holds 4,096 × 128,256 ≈ 525 million weights: softmax regression with 128,256 classes.

Temperature softmax as used in knowledge distillation, Hinton, Vinyals & Dean, Distilling the Knowledge in a Neural Network (2015), where a high T exposes the teacher's "dark knowledge" in the small probabilities of wrong classes. The same knob controls randomness in LLM sampling.

Training a softmax classifier

Encode labels as integers (or one-hot); stratified train/test split.
Standardise features; add engineered features if classes have curved shapes.
Compute logits t = Wx + b, then the loss with a fused, stable log-softmax + cross-entropy (never take the log of a softmax you computed separately).
Optimise with L-BFGS (small data) or SGD/Adam (large); add L2 regularization (it also removes the redundancy in the logits).
Evaluate with per-class precision and recall, a confusion matrix, macro-F1 for imbalanced classes, and log-loss for probability quality.
For decisions, use argmax, or thresholds and costs per class; check calibration if probabilities feed decisions.

Use it, or reach for something else?

Use softmax regression when

Classes are mutually exclusive and you need coherent probabilities
You want a fast, interpretable multi-class baseline (e.g. text with TF-IDF)
You're building the output layer of any neural classifier

Use something else when

Items can carry several labels → one sigmoid per label
Classes are ordered (ratings 1–5) → ordinal regression
Boundaries are strongly non-linear in tabular data → gradient boosting

The answer that sounds right

The trap

"Divide each score by the total"

The first idea: just normalise, p_k = t_k / Σ t_j. Scores of (2, 1, 1) give (0.5, 0.25, 0.25). Looks fine, until a score is negative or the scores sum to zero:

Banana's score

Apple and orange scores fixed at 2.0 and 1.0. Plain division produces negative "probabilities", values above 1, or a division by zero as banana's score moves. Softmax always gives a valid distribution, and keeps the order of the scores.

What breaks

Plain normalisation needs every score to be positive. Exponentiating first fixes that (e^t > 0 always), and it has a deeper reason: it's what you get by extending the logistic model's log-odds idea to K classes.

What to keep

The lesson: softmax regression is logistic regression for K classes: one logit per class, exponentiate and normalise, cross-entropy loss, gradient φ − y. It's linear, convex, and the output layer of almost every modern classifier and language model.

Mutually exclusive classes only. If an item can have several labels at once (multi-label tagging), use independent sigmoids per label, not one softmax.
Softmax probabilities are not automatically calibrated, especially in big neural nets, which tend to be overconfident. Temperature scaling on a validation set is the usual fix. my take
Huge K is expensive. The normaliser sums over every class; with very large vocabularies, training uses tricks like sampled softmax, and inference often only needs the top few logits.
Linear boundaries, just like logistic regression.

Say it

What is softmax regression, and how is it trained?

Say your answer out loud

Aim for 30 seconds. It feels different from reading it, and that's the point.

Tick what you covered

It's multinomial logistic regression, with one weight vector and logit per class.
Softmax exponentiates and normalises the logits into a distribution.
It's trained by minimising cross-entropy.
The gradient on the logits is predicted probabilities minus the one-hot label.
It's convex, and it reduces to the sigmoid with two classes.
It's the output layer of neural classifiers and LLMs.