Softmax for output layers: Exercises — AI-ML deep dive

Before we begin, a one-line refresher of the only machine we use, so no symbol is unearned:

We also reuse two companions from the parent:

The cross-entropy loss for a one-hot target $t$ (a vector that is $1$ on the true class, $0$ elsewhere): $L = - \sum_{i} t_{i} lo g p_{i}$ . See Cross-Entropy Loss.
Its clean gradient w.r.t. logits: $\frac{\partial L}{\partial z _{k}} = p_{k} - t_{k}$ .

Let's climb.

Level 1 — Recognition

Exercise 1.1

Which of these vectors could be the output of a softmax (over 3 classes)? (a) $[0.2, 0.5, 0.3]$ (b) $[0.6, 0.6, - 0.2]$ (c) $[0.1, 0.1, 0.1]$ (d) $[1.0, 0.0, 0.0]$

Recall Solution

A softmax output must obey two laws: every entry lies in $(0, 1)$ , and the entries sum to exactly 1.

(a) sums to $0.2 + 0.5 + 0.3 = 1.0$ , all positive → valid. ✔
(b) contains $- 0.2 < 0$ → impossible, $e^{z} > 0$ can never give a negative → invalid. ✘
(c) sums to $0.3 \neq = 1$ → invalid (it forgot to normalize). ✘
(d) sums to $1$ , but two entries are exactly $0$ and one is exactly $1$ . Softmax outputs are strictly in $(0, 1)$ — it can approach but never reach $0$ or $1$ (that would need an infinite logit gap). So this is a limiting case, not a true softmax output. ✘

Answer: only (a).

Exercise 1.2

The softmax of $z = [5, 5, 5]$ is what vector, without computing any exponentials?

Recall Solution

All logits are equal, so all $e^{z_{i}}$ are equal, so all shares are equal. With $K = 3$ equal shares summing to $1$ , each is $1/3$ . Answer: $[\frac{1}{3}, \frac{1}{3}, \frac{1}{3}] \approx [0.333, 0.333, 0.333]$ . Equal scores → uniform distribution, regardless of the common value.

Level 2 — Application

Exercise 2.1

Compute softmax of $z = [0, ln 2, ln 3]$ exactly (as fractions).

Recall Solution

The clever part: $e^{l n x} = x$ , so exponentiating undoes the log.

Tops: $e^{0} = 1$ , $e^{l n 2} = 2$ , $e^{l n 3} = 3$ .
Sum (normalizer): $1 + 2 + 3 = 6$ .
Divide: $[\frac{1}{6}, \frac{2}{6}, \frac{3}{6}] = [\frac{1}{6}, \frac{1}{3}, \frac{1}{2}]$ .

Answer: $[1/6, 1/3, 1/2] \approx [0.167, 0.333, 0.500]$ . The logits were literally chosen as logs of the desired ratio $1 : 2 : 3$ .

Exercise 2.2

Take logits $z = [2.0, 1.0, 0.1]$ . First compute their softmax from scratch on this page. Then add $10$ to every logit, giving $[12.0, 11.0, 10.1]$ , and find the new softmax.

Recall Solution

Compute the base softmax first — no numbers borrowed.

Exponentiate: $e^{2.0} = 7.389$ , $e^{1.0} = 2.718$ , $e^{0.1} = 1.105$ .
Sum (normalizer): $7.389 + 2.718 + 1.105 = 11.212$ .
Divide: $[\frac{7.389}{11.212}, \frac{2.718}{11.212}, \frac{1.105}{11.212}] = [0.659, 0.242, 0.099]$ .

Now the shift. Shift-invariance: adding the same constant $c$ to every logit multiplies every top by $e^{c}$ , top and bottom, so it cancels: $\frac{e ^{z_{i} + c}}{\sum _{j} e ^{z_{j} + c}} = \frac{e ^{c} e ^{z_{i}}}{e ^{c} \sum _{j} e ^{z_{j}}} = \frac{e ^{z_{i}}}{\sum _{j} e ^{z_{j}}} .$ Answer: the shifted softmax is unchanged, $[0.659, 0.242, 0.099]$ . Softmax only cares about differences between logits, never their absolute level.

Exercise 2.3

For binary logits $z = [z_{1}, z_{2}]$ , show softmax of class 1 equals the sigmoid of the gap $z_{1} - z_{2}$ , where the sigmoid is $σ (x) = \frac{1}{1 + e ^{- x}}$ .

Recall Solution

Start from softmax and divide top and bottom by $e^{z_{1}}$ to "measure everything relative to class 1": $p_{1} = \frac{e ^{z_{1}}}{e ^{z_{1}} + e ^{z_{2}}} = \frac{1}{1 + e ^{z_{2} - z_{1}}} = \frac{1}{1 + e ^{- (z_{1} - z_{2})}} = σ (z_{1} - z_{2}) .$ Answer: $p_{1} = σ (z_{1} - z_{2})$ . This is why 2-class softmax and the Sigmoid Activation are the same machine — the sigmoid is just softmax that has already been reduced to a single logit gap. See also Logits and Log-Odds.

Level 3 — Analysis

Before the next exercise we must earn one new symbol, the Kronecker delta $δ_{ik}$ .

Exercise 3.1

Verify the derivative identity $\frac{\partial p _{i}}{\partial z _{k}} = p_{i} (δ_{ik} - p_{k})$ for the diagonal case $i = k$ , working from $p_{i} = e^{z_{i}} / S$ with $S = \sum_{j} e^{z_{j}}$ .

Recall Solution

$p_{i} = e^{z_{i}} / S$ is a ratio, so we use the quotient rule $(\frac{u}{v})^{'} = \frac{u ^{'} v - u v ^{'}}{v ^{2}}$ . We differentiate w.r.t. $z_{i}$ (the diagonal, $k = i$ ):

Top $u = e^{z_{i}}$ , so $\partial u / \partial z_{i} = e^{z_{i}}$ .
Bottom $v = S$ . Since $S$ contains the term $e^{z_{i}}$ , $\partial S / \partial z_{i} = e^{z_{i}}$ .

Plug in: $\frac{\partial p _{i}}{\partial z _{i}} = \frac{e ^{z_{i}} S - e ^{z_{i}} e ^{z_{i}}}{S ^{2}} = \frac{e ^{z_{i}}}{S} \cdot \frac{S - e ^{z_{i}}}{S} = p_{i} (1 - p_{i}) .$ On the diagonal $i = k$ , so $δ_{ik} = δ_{ii} = 1$ , and this is exactly $p_{i} (δ_{ii} - p_{i})$ . ✔ (Off-diagonal $k \neq = i$ : only the bottom depends on $z_{k}$ , giving $\frac{- e ^{z_{i}} e ^{z_{k}}}{S ^{2}} = - p_{i} p_{k}$ . Here $δ_{ik} = 0$ , so $p_{i} (0 - p_{k}) = - p_{i} p_{k}$ matches.)

Exercise 3.2

A network predicts $p = [0.7, 0.2, 0.1]$ and the true class is index 1 (0-based), so $t = [0, 1, 0]$ . Compute the loss gradient $\partial L / \partial z_{k}$ for each $k$ , and say which logit gets pushed up.

Recall Solution

The magic collapse: $\frac{\partial L}{\partial z _{k}} = p_{k} - t_{k}$ (prediction minus target).

$k = 0$ : $0.7 - 0 = + 0.7$
$k = 1$ : $0.2 - 1 = - 0.8$
$k = 2$ : $0.1 - 0 = + 0.1$

Gradient descent moves each logit opposite to its gradient. A negative gradient means the logit is pushed up. Only $k = 1$ (the true class) has a negative gradient, so logit 1 is pushed up; the two wrong logits are pushed down. Exactly what learning should do — see Backpropagation. Answer: $[+ 0.7, - 0.8, + 0.1]$ ; logit 1 (true class) goes up.

Exercise 3.3

Show the sum of all cross-entropy gradients over $k$ is zero: $\sum_{k} \frac{\partial L}{\partial z _{k}} = 0$ . Why does this make sense geometrically?

Recall Solution

$\sum_{k} (p_{k} - t_{k}) = \sum_{k} p_{k} - \sum_{k} t_{k} = 1 - 1 = 0,$ because probabilities sum to $1$ and a one-hot target also sums to $1$ . Meaning: the gradient can never shift all logits in the same direction — that would be the shift from Ex 2.2, which changes nothing. So training only ever adjusts the relative pushes: every unit pushed up somewhere is exactly balanced by pushes down elsewhere.

Level 4 — Synthesis

Exercise 4.1 — Temperature

Define tempered softmax $p_{i} (T) = \frac{e ^{z_{i} / T}}{\sum _{j} e ^{z_{j} / T}}$ for temperature $T > 0$ . For $z = [2, 1, 0]$ , compute the winning probability $p_{1}$ at $T = 1$ and at $T = 2$ , and state the limits as $T \to 0^{+}$ and $T \to \infty$ .

Recall Solution

Dividing logits by $T$ shrinks (large $T$ ) or stretches (small $T$ ) the gaps.

$T = 1$ : tops $e^{2}, e^{1}, e^{0} = 7.389, 2.718, 1$ ; sum $11.107$ ; $p_{1} = 7.389/11.107 = 0.665$ .
$T = 2$ : logits become $[1, 0.5, 0]$ ; tops $e^{1}, e^{0.5}, e^{0} = 2.718, 1.649, 1$ ; sum $5.367$ ; $p_{1} = 2.718/5.367 = 0.506$ .

Bigger $T$ softens (0.665 → 0.506, closer to uniform $1/3$ ).

$T \to \infty$ : all gaps $\to 0$ , so $p \to [\frac{1}{3}, \frac{1}{3}, \frac{1}{3}]$ (uniform).
$T \to 0^{+}$ : gaps blow up, the largest logit dominates completely, so $p \to [1, 0, 0]$ (hard argmax).

Answers: $p_{1} (1) = 0.665$ , $p_{1} (2) = 0.506$ ; limits uniform ( $T \to \infty$ ) and one-hot argmax ( $T \to 0^{+}$ ). This is Temperature Scaling.

The figure below traces this: the horizontal axis is temperature $T$ , the vertical axis is softmax probability. Three curves — burnt-orange for class 1 ( $z = 2$ ), teal for class 2 ( $z = 1$ ), plum for class 3 ( $z = 0$ ). Read it right-to-left: as $T$ grows the three curves squeeze toward the dashed uniform line at $1/3$ (all classes tie); as $T$ shrinks toward $0$ the orange curve rockets to $1$ while the others fall to $0$ (argmax). The two grey verticals mark the $T = 1$ and $T = 2$ answers you just computed.

Exercise 4.2 — Maximum entropy

Softmax is the distribution of maximum entropy $H (p) = - \sum_{i} p_{i} lo g p_{i}$ given fixed expected logit. Without heavy calculus, argue why the uniform distribution (Ex 1.2's case) is the special softmax that maximizes entropy outright.

Recall Solution

Entropy $H$ measures "spread-outness / uncertainty." It is largest when no outcome is favored — i.e. all $p_{i}$ equal. When the logits are all equal, softmax returns exactly the uniform $[\frac{1}{K}, \dots, \frac{1}{K}]$ , whose entropy is the maximum possible value $lo g K$ . Any unequal logits create a preference, lowering entropy. So softmax with equal logits = the maximum-entropy point, the "least-biased guess." For the full derivation see Maximum Entropy Distributions. Check for $K = 3$ : $H = - 3 \cdot \frac{1}{3} ln \frac{1}{3} = ln 3 \approx 1.099$ .

Level 5 — Mastery

Exercise 5.1 — The double-softmax destruction

You accidentally apply softmax twice: $q = softmax (softmax (z))$ . Starting from $z = [2, 1, 0]$ , compute $q$ and compare to the correct single softmax $p$ . What went wrong?

Recall Solution

First softmax: $p = [0.665, 0.245, 0.090]$ (Ex 4.1, $T = 1$ ).
Second softmax treats these tiny numbers as new logits. Their gaps are now minuscule ( $0.665 - 0.090 = 0.575$ instead of $2$ ), so exponentiating barely separates them: tops $e^{0.665}, e^{0.245}, e^{0.090} = 1.945, 1.278, 1.094$ ; sum $4.317$ ; $q = [0.451, 0.296, 0.253]$ .

Answer: $q \approx [0.451, 0.296, 0.253]$ — much flatter than $p$ . Applying softmax again shrinks the gaps and pushes everything toward uniform, destroying information. Rule: apply softmax exactly once, to raw logits.

Exercise 5.2 — Overflow-safe computation

Compute softmax of $z = [1000, 1000, 1002]$ by hand, safely.

Recall Solution

Naïve $e^{1002}$ overflows to inf. Use shift-invariance (Ex 2.2): subtract $max = 1002$ .

Shifted: $[- 2, - 2, 0]$ .
Tops: $e^{- 2}, e^{- 2}, e^{0} = 0.1353, 0.1353, 1$ ; sum $= 1.2707$ .
Divide: $[0.1065, 0.1065, 0.7870]$ .

Answer: $\approx [0.1065, 0.1065, 0.787]$ , computed with all exponents $\leq 0$ — no overflow, and identical to the un-shifted result.

Exercise 5.3 — Degenerate / limiting input

What does softmax return for $z = [c, c, \dots, c]$ (all equal, any $c$ ), and what happens to $p_{1}$ as $z_{1} \to + \infty$ with the other logits held fixed? Cover both edges.

Recall Solution

Equal case: every top equals $e^{c}$ , so the answer is uniform $[\frac{1}{K}, \dots, \frac{1}{K}]$ — independent of $c$ (this combines Ex 1.2 and shift-invariance). Runaway case ( $z_{1} \to + \infty$ ): $e^{z_{1}}$ dominates the sum, so $p_{1} = \frac{e ^{z_{1}}}{e ^{z_{1}} + (fixed)} \to 1, p_{i \neq = 1} = \frac{e ^{z_{i}}}{e ^{z_{1}} + \dots} \to 0.$ Answer: uniform for equal logits; approaches one-hot $[1, 0, \dots, 0]$ as one logit runs to $+ \infty$ . Note $p_{1}$ never exactly equals 1 for finite logits — it only limits there, confirming outputs live in the open interval $(0, 1)$ (why Ex 1.1(d) failed).

The figure below makes the runaway case visible. It fixes $z_{2} = z_{3} = 0$ and slides $z_{1}$ up through the values $0, 2, 4, 8$ ; each cluster of three bars is one softmax. The horizontal axis groups the three classes, the vertical axis is probability. Watch the class-1 bar climb toward the dashed line at $1$ while the class-2 and class-3 bars collapse toward $0$ — exactly the $p_{1} \to 1$ , $p_{i \neq = 1} \to 0$ limit. The annotation flags that the bar approaches but never reaches $1$ , reconfirming why Ex 1.1(d)'s exact $[1, 0, 0]$ is not a true softmax output.

Active Recall

Level up recognized

L1→L5 = Recognition, Application, Analysis, Synthesis, Mastery.

Softmax of all-equal logits

uniform

[1/ K, \dots, 1/ K]

, independent of the common value.

Safe-compute trick

subtract

max_{k} z_{k}

before exponentiating (exact, prevents overflow).

Loss gradient w.r.t. logit

\partial L / \partial z_{k} = p_{k} - t_{k}

Kronecker delta meaning

δ_{ik} = 1

i = k

, else

0

— a same-index switch.

Connections

Cross-Entropy Loss — supplies the $- \sum t_{i} lo g p_{i}$ that yields the clean $p - t$ gradient (Ex 3.2, 3.3).
Sigmoid Activation — the 2-class reduction of softmax (Ex 2.3).
Logits and Log-Odds — the raw inputs; gaps are log-odds.
Maximum Entropy Distributions — justifies the exponential form (Ex 4.2).
Backpropagation — where the $p - t$ gradient flows.
Temperature Scaling — the $T$ knob of Ex 4.1 / 5-limits.