1.7 — From binary to multi-class: the softmax

1. Why this session exists

The sigmoid (1.2) handled $y \in \{0, 1\}$. For $y \in \{1, \dots, K\}$ we need a probability-simplex output. The softmax is the canonical generalisation. It will play four roles:

• The natural multinomial extension of the sigmoid (this session).
• The MLE of a multinomial-logit observation model (this session).
• The output head of essentially every modern classifier and language model (Chapters 5–8).
• The gradient sink with the cleanest backward pass in all of deep learning, $\partial L / \partial z = \hat y - y$, generalising 1.2 — the topic of 1.8.[1]

2. Setup

Labels: $y \in \{1, \dots, K\}$. We will mostly work with the one-hot encoding,

Per-class linear scores ("logits"):

stacked into a vector $z \in \mathbb{R}^K$. The weight matrix is $W \in \mathbb{R}^{K \times n}$ with rows $w_k^\top$; the bias is $b \in \mathbb{R}^K$.

3. The softmax

Define the softmax as

Three properties.

Probability simplex. Each entry is in $(0, 1)$, and the entries sum to $1$:

So $\hat y := \operatorname{softmax}(z)$ is a valid distribution over the $K$ classes.

Shift invariance. From 0.10 §4, $\operatorname{softmax}(z + c\,\mathbf{1}) = \operatorname{softmax}(z)$ for any scalar $c$. We use this to stabilise the computation by subtracting $\max_j z_j$.

Reduction to sigmoid. For $K = 2$ with logits $(z_1, z_2)$,

So the two-class softmax depends only on the difference of logits, which is the sigmoid of the difference. Setting $z_2 \equiv 0$ recovers the binary sigmoid; the binary case is over-parameterised, which is why for $K = 2$ we usually use one logit and $\sigma$.

4. The multinomial-logit observation model

Postulate that, conditional on $x$,

This is the multinomial-logit model — the multi-class generalisation of the Bernoulli logistic model in 1.2 (which collapses out of $K = 2$).

The data likelihood for a single example $(x, y)$ with one-hot label $y$:

i.e. just the predicted probability of the correct class — the product is $1$ for every $k$ except the true class.

5. MLE = cross-entropy

The log-likelihood for $m$ examples is

Negating and averaging gives the multi-class cross-entropy training loss:

The same construction as 1.2 (BCE) and 1.1 (least squares): negative average log-likelihood of an observation model. For one-hot $y$ the inner sum reduces to picking out the predicted probability of the true class, $\log \hat y_{i,\, y_i}$.

6. Cross-entropy on logits via log-sum-exp

Inserting the definition of softmax,

The right-hand side is the log-sum-exp expression from 0.10 §4. Two consequences.

Stability for free. Pair this loss with the stable log-sum-exp implementation $\max_j z_j + \log \sum_j e^{z_j - \max_j z_j}$ and you get a fused softmax + cross-entropy that never overflows.

Direct link to information theory. $-\log p(y \mid x)$ is the information content of observing $y$ under the model. Summed over the dataset and divided by $m$, it is the empirical cross-entropy between the data distribution and the model (0.6, 0.7). Minimising it is the Cover-and-Thomas "code-length minimisation" reading of MLE.

7. Decision rule and Bayes optimality

At inference, the classifier predicts $\hat y = \arg\max_k \hat y_k = \arg\max_k z_k$ (since $\exp$ is monotone). If the multinomial-logit model is correctly specified, this rule coincides with the Bayes-optimal classifier that minimises the misclassification rate.

In practice the model is misspecified (the world is not log-linear in the raw features), but the softmax + MLE pairing remains an excellent default — it is convex in $(W, b)$, scales linearly with $K$, and gives calibrated probabilities under mild assumptions.

8. Takeaways

• Softmax: smooth, simplex-valued, shift-invariant, reduces to sigmoid for $K = 2$.
• Multinomial-logit model: $p(y = k \mid x) = \operatorname{softmax}(Wx + b)_k$.
• MLE of this model = multi-class cross-entropy loss.
• Combined with log-sum-exp, the loss admits a stable fused computation directly on logits.