Loss Functions & Output Units

A network's hidden layers learn a representation $h=f_{\theta }(x)$; the output unit is the thin final transform that turns $h$ into a prediction, and the loss scores that prediction against the target. These two are not chosen independently. Fix the conditional distribution $p_{\theta }(y\mid x)$ you intend the network to model, and maximum likelihood determines the output activation and the loss as a matched pair. The whole lesson is one table and the derivations behind its five rows.

The pattern is rigid: pick the distribution that matches the target's type — real, binary, one-of-$K$, several-of-$K$, a nonnegative integer — let its natural parameter be the network's pre-activation, and the loss is the negative log-likelihood of that distribution. The rest of the lesson derives each row.¹

Maximum likelihood is the loss

Every row above is one instance of a single principle. The network outputs the parameters of a conditional distribution $p_{\theta }(y\mid x)$, and training maximizes the likelihood of the observed targets, equivalently minimizes the negative log-likelihood (NLL).

This single choice removes the guesswork from loss design: the loss is read off as $-\log p_{\theta }(y\mid x)$ for whichever $p_{\theta }$ suits the target. Two consequences recur throughout.²

• The output activation is whatever maps an unbounded score into a valid parameter of $p_{\theta }$ — a probability needs $(0,1)$, so a sigmoid; a rate needs $(0,\infty )$, so a softplus.
• The $\log$ in the NLL undoes the $\exp$ in every exponential-family output unit, and that cancellation is why the gradients below collapse so cleanly. A loss without that $\log$ — squared error on a sigmoid — leaves the $\exp$ intact and the gradient saturates.

The second consequence is worth making precise, because it is the mechanism behind every clean gradient in this lesson. An exponential-family output writes the model probability as $p_{\theta }=e^{a(z)}/Z(z)$ for some score $z$ and normalizer $Z$. Taking the negative log turns the ratio into a difference,

so the $\exp$ that built the probability is gone and the loss is linear in the score plus a smooth log-normalizer. Differentiating a linear term gives a constant, and differentiating $\log Z$ returns the model's own predicted probabilities. The two pieces combine into $\hat{y}-y$ with no leftover $\exp$ and no activation derivative multiplying the signal. Drop the $\log$ — square the residual of a sigmoid instead — and the $\exp$ survives inside the gradient as the factor $\hat{y}(1-\hat{y})$, which vanishes exactly where the model is most wrong. The last section of this lesson shows that failure numerically.

Regression: linear output, squared error

For a real-valued target take $p_{\theta }(y\mid x)=\mathcal{N}\left(y;\hat{y},{\sigma }^2\right)$ with the mean predicted by an identity output, $\hat{y}=w^{\top }h+b$, and ${\sigma }^2$ fixed. The NLL is

Drop the additive constant and the scale, and the loss is the mean squared error $\frac{1}{2}(\hat{y}-y{)}^2$. Squared error is not an arbitrary penalty but the log-likelihood of a Gaussian with fixed variance. Its gradient with respect to the score is the residual, with no saturating factor:

Because the output is the identity, this residual passes straight to the weights as $(\hat{y}-y)h$. The lesson's recurring theme — gradient equals prediction minus target — starts here, in its simplest form.

Softmax: the multiclass output unit

For a one-of-$K$ label, the network emits a vector of $K$ scores (logits) $z\in {\mathbb{R}}^K$, and the softmax turns them into a categorical distribution over the classes.

The figure traces the three stages (score, exponentiate, normalize) that carry a logit vector to a probability distribution.

Softmax inherits two structural facts. It is shift-invariant: adding a constant $c$ to every logit multiplies numerator and denominator by $e^c$ and leaves $\hat{y}$ unchanged. And the raw form overflows: $e^{z_i}$ explodes for large logits. Both are fixed at once by subtracting the maximum logit before exponentiating, which is the form every library actually computes.

A temperature $T$ rescales the logits before the softmax, $\mathrm{softmax}(z/T)$, interpolating between a one-hot argmax ($T\to 0$) and the uniform distribution ($T\to \infty$). The bar chart shows the same logits sharpened and flattened.

The softmax + cross-entropy gradient

The matched loss for a softmax output is the categorical cross-entropy. With a one-hot target $y$ (so $y_c=1$ for the true class $c$ and $0$ otherwise),

The result that makes this pairing canonical is that the gradient with respect to the logits collapses to the residual $\hat{y}-y$. Derive it by differentiating $\ell =-\log {\hat{y}}_c$ through the softmax. First the log-sum-exp:

The first term is the indicator $[k=c]$, which is precisely $y_k$ for a one-hot target; the second term is ${\hat{y}}_k$ by the softmax definition. Therefore, for every coordinate,

This is the central fact of classification training. The error signal passed back to the network is just predicted distribution minus true distribution, with no activation derivative attenuating it.³

A worked pass through the numbers

For example, take $K=3$ classes with true class $c=1$ (so the one-hot target is $y=(1,0,0)$), and logits $z=(2,1,0)\in {\mathbb{R}}^3$. The stable softmax subtracts ${\max }_k{z}_k=2$:

Dividing gives the predicted distribution $\hat{y}=(0.665,0.245,0.090)$, which sums to $1$ and is identical to what the naive formula returns. The per-example loss is $\ell =-\log {\hat{y}}_1=-\log 0.665=0.408$ nats. The gradient with respect to the logits is one subtraction, $\hat{y}-y$:

The true-class logit gets a negative gradient, so gradient descent raises $z_1$; the two wrong-class logits get positive gradients, so descent lowers them. The magnitudes sum to zero ($-0.335+0.245+0.090=0$), a direct consequence of ${\sum }_i({\hat{y}}_i-y_i)={\sum }_i{\hat{y}}_i-{\sum }_i{y}_i=1-1=0$: softmax cross-entropy only ever redistributes probability mass, never creates or destroys it. In a minibatch of $B$ examples the logits are a matrix $Z\in {\mathbb{R}}^{B\times K}$, the softmax acts row-wise, and the gradient $\hat{Y}-Y\in {\mathbb{R}}^{B\times K}$ has the same shape as the logits, ready to seed the backward pass.

Binary: sigmoid + binary cross-entropy

The two-class case is softmax with $K=2$ collapsed to a single logit. Model $p_{\theta }(y=1\mid x)=\hat{y}=\sigma (z)$ with the logistic sigmoid$\sigma (z)=1/(1+e^{-z})$, take $y\in \{0,1\}$, and the Bernoulli NLL is the binary cross-entropy

Using ${\sigma }^{\prime }(z)=\sigma (z)(1-\sigma (z))=\hat{y}(1-\hat{y})$, the awkward $1/[\hat{y}(1-\hat{y})]$ from the log derivative is cancelled by the sigmoid's own derivative, leaving the same residual:

This is the binary specialization of the softmax result, and it is why both output units share one backward implementation: combine the activation and the loss into a single op whose gradient is $\hat{y}-y$, never materializing the unstable intermediate $1/[\hat{y}(1-\hat{y})]$.

Why squared error is wrong for a classifier

A tempting mistake is to keep MSE but put it on a sigmoid output, ${\ell }_{\text{MSE}}=\frac{1}{2}(\hat{y}-y{)}^2$ with $\hat{y}=\sigma (z)$. The problem is the gradient with respect to the score. By the chain rule it picks up the sigmoid's derivative as a multiplicative factor:

The factor $\hat{y}(1-\hat{y})$ causes saturation. When the model is confidently wrong — say $y=1$ but $z\ll 0$, so $\hat{y}\approx 0$ — the residual $\hat{y}-y\approx -1$ is large, yet $\hat{y}(1-\hat{y})\approx 0$ drives the MSE gradient to nearly zero. The example that most needs correcting produces almost no learning signal. Cross-entropy has no such factor, so its gradient stays order $1$ exactly where it matters. The table below quantifies the failure.⁴

The first row is the key case: a $150\times$ weaker gradient on the worst prediction. Two figures plot the comparison: first the per-example loss as a function of the predicted probability of the true class, then the gradient magnitude across the score.

The BCE curve approaches its maximum gradient exactly in the region where the model is most wrong ($z\to -\infty$); the MSE curve collapses to zero at both extremes, learning slowly precisely when it should learn fastest.

The output-unit map

Every row of the master table is the same construction applied to a different target type. The branching diagram fixes the mapping from what kind of thing $y$ is to the output unit and loss that maximum likelihood forces.

Two rows of the table extend the binary and count cases.

Multi-label. When an example may carry several labels at once (an image tagged both beach and sunset), the classes are not mutually exclusive, so softmax is wrong: it forces ${\sum }_i{\hat{y}}_i=1$. Instead apply an independent sigmoid per class and sum $K$ binary cross-entropies,

which is the NLL of $K$ independent Bernoullis. Each logit's gradient is its own residual ${\hat{y}}_i-y_i$.

Count. For a nonnegative integer target (events per interval) model $p_{\theta }(y\mid x)=\text{Poisson}(\lambda )$ with rate $\lambda =\mathrm{softplus}(z)=\log (1+e^z)>0$. The NLL drops the data-only $\log y!$ term:

Softplus keeps the rate strictly positive while staying smooth, where a bare $e^z$ would overflow and a ReLU would kill the gradient at $z<0$.

When the network predicts its own uncertainty

The regression row fixed ${\sigma }^2$ and threw it away. Nothing forces that. If the target's noise varies across inputs (quiet in one region of $x$, noisy in another), the network can predict the variance too. This is a heteroscedastic Gaussian output: two heads, $\hat{y}(x)=w_{\mu }^{\top }h+b_{\mu }$ for the mean and ${\sigma }^2(x)=\mathrm{softplus}(w_{\sigma }^{\top }h+b_{\sigma })>0$ for the variance. The NLL keeps both terms this time:

The first term is a squared error weighted by the network's own confidence: where the model predicts small ${\sigma }^2$ it is penalized heavily for being wrong, and where it predicts large ${\sigma }^2$ the residual is discounted. The second term, $\frac{1}{2}\log {\sigma }^2$, penalizes that discount, and it blocks the trivial solution of setting ${\sigma }^2\to \infty$ to zero out the first term. The two terms balance at the true noise level, so the network learns a calibrated variance alongside the mean.

A mixture-density network (MDN) goes further when the target is genuinely multimodal, one $x$ that maps to several plausible $y$ where any single Gaussian would average them into an implausible midpoint. The output parameterizes a mixture of $M$ Gaussians, emitting mixing weights ${\pi }_m(x)$ through a softmax over $M$ components, plus a mean ${\mu }_m(x)$ and variance ${\sigma }_m^2(x)$ per component:

The loss is still a negative log-likelihood, but now of a mixture, so the log sits outside the sum and does not collapse into a clean residual. The principle is unchanged though: pick the conditional distribution that matches the target, and the loss is $-\log p_{\theta }(y\mid x)$.³

Label smoothing and class imbalance

Two practical edits to the target keep cross-entropy well-behaved. Both modify $y$, not the loss.

Label smoothing. A hard one-hot target requires ${\hat{y}}_c=1$, which is reachable only as $z_c\to \infty$ — the logits grow without bound and the model becomes overconfident. Label smoothing replaces the one-hot target with a softened version that mixes in the uniform distribution:

for a small $\epsilon$ (commonly $0.1$). The cross-entropy minimum now sits at finite logits, regularizing confidence and improving calibration.⁵

Class imbalance. When one class dominates the data, the majority dominates the averaged loss and the minority's gradient is negligible. The standard fix reweights each class's contribution by a factor ${\alpha }_y$ inversely related to its frequency:

with $n_k$ the count of class $k$. This restores per-class influence without touching the architecture; the gradient is simply scaled per example by ${\alpha }_y$.

Losses beyond maximum likelihood

Maximum likelihood is the base construction, but three widely used losses extend it in ways Goodfellow predates or treats only in passing.

Focal loss down-weights the easy examples. Class reweighting scales by label frequency, but in extreme imbalance (dense object detection, where background boxes outnumber objects thousands to one) even the correctly classified easy negatives dominate the summed gradient. Lin et al.'s focal loss (2017) multiplies the cross-entropy of each example by $(1-{\hat{y}}_c{)}^{\gamma }$, so a confident correct prediction (${\hat{y}}_c\to 1$) contributes almost nothing and the loss concentrates on the hard, still-misclassified cases. It is cross-entropy with a per-example confidence gate, and it is the modern default wherever imbalance is severe.

Calibration is a separate axis from accuracy. The clean $\hat{y}-y$ gradient trains a classifier to be accurate, but a network can be accurate and still badly calibrated — its softmax probabilities systematically over- or under-state the true likelihood. Guo et al. (2017) documented that modern deep nets are overconfident and that a one-parameter fix, temperature scaling — dividing the logits by a scalar $T$ tuned on held-out data, exactly the temperature of the softmax section — recovers calibration without touching accuracy. Label smoothing above is the training-time cousin of the same concern. The lesson from the literature is that the loss you minimize and the probabilities you can trust are related but not identical.

Contrastive losses are structurally cross-entropy. Self-supervised learning trains without labels by a loss that is structurally the softmax cross-entropy of this lesson. The InfoNCE objective (van den Oord et al., 2018; SimCLR, Chen et al., 2020) treats one positive pair against $N$ negatives as an $(N+1)$-way classification problem: the logits are similarities, the target is the positive is class $0$, and the loss is $-\log$ of the softmax over similarities. The residual gradient $\hat{y}-y$ derived here drives it too — the same output unit doing a different job.⁶

The shared backward op

Across four of the five rows the backward pass is identical: the gradient at the score is the residual $\hat{y}-y$. That is the practical payoff of letting maximum likelihood choose the pairing: the forward unit and the loss are designed together so their derivatives cancel into one subtraction.

The lesson reduces to one instruction: choose the output distribution that matches the target's type, and the activation, the loss, and this clean residual gradient all follow. The matched pair then plugs straight into backpropagation as the seed of the backward pass, and into the optimizer as the gradient it descends.