Output Layers and Loss Pairings

The final layer and the loss are one decision, not two. The output layer defines what numbers the network emits; the loss defines what those numbers mean. When they disagree, training still runs — gradients are computed, the loss goes down, and the model quietly optimizes the wrong objective. So start from the target tensor and work backwards, never from a favorite layer.

Two columns in that table surprise people. First, logits — the raw pre-activation values. Modern PyTorch classification losses take logits and fold the sigmoid or softmax inside a numerically stable fused kernel, so your model’s forward pass ends at the Linear layer with no activation before the loss. When you need probabilities for reporting, thresholds, or calibration, apply torch.sigmoid(logits) or torch.softmax(logits, dim=1) after the fact. And since softmax is monotone, argmax over logits equals argmax over its output — accuracy needs no activation either.

Three questions walk the table and you can answer them about any new request in under a minute: is the target a number or a class; are the classes mutually exclusive; is there one label per example or many. A bounded number takes a squashing output, an unbounded number takes a bare linear one, classes take a K-way head with a log-likelihood from the cross-entropy family, and counts take PoissonNLL, derived from the Poisson distribution’s likelihood rather than a distance. Decide the pairing on paper before the first epoch — in production a wrong pairing trains happily and only shows up as a metric that sits slightly, inexplicably off.

Second, the target formats. CrossEntropyLoss expects a LongTensor of class indices, shape (N,), not one-hot rows — hand it float targets and you get an error about the expected scalar type Long, or in newer versions a silent reinterpretation as soft labels you did not intend. BCEWithLogitsLoss expects float targets in {0, 1} with the same shape as the logits; a (N,) versus (N, 1) mismatch triggers broadcasting that quietly changes the loss rather than failing. The multiclass versus multi-label split is decided by the loss, not the architecture: softmax forces the K outputs to compete for a total probability of 1, which encodes mutually exclusive classes, while sigmoid scores each output independently, which is what lets an article be both sports and finance.

For regression, two pairing notes earn their keep in practice. An unbounded target needs a linear output — put a sigmoid on a price model and you have capped every prediction at 1, a bug that looks like terrible training rather than a wrong layer. Standardize targets to zero mean and unit variance before MSE, or the initial loss starts in the millions and the first updates are dominated by scale instead of signal. And when the data has outliers, MSE squares them into dominating gradients; Activation Functions covered why squashing units saturate, and here the analogous fix is HuberLoss — quadratic near zero, linear in the tails, so a handful of wild points cannot hijack a batch.

For K classes with logits z and the true class index y, categorical cross-entropy is the negative log probability of the true class:

L = −log p_y   where   p = softmax(z)
expanded:  L = −z_y + log Σ_j exp(z_j)

Work it on real numbers: logits (2.0, 1.0, 0.1) with true class 0. exp(2) = 7.389, exp(1) = 2.718, exp(0.1) = 1.105, summing to 11.212, so p = (0.659, 0.242, 0.099) and L = −log(0.659) = 0.417. Sanity checks worth memorizing: perfectly separated logits give p_y ≈ 1 and L ≈ 0, while identical logits give a uniform distribution and L = log K (log 3 ≈ 1.10 for three classes). The loss punishes confidence that is wrong, and only confidence.

Now the reason every serious implementation looks different from the formula. The naive version overflows:

def naive_ce(logits, y):                     # do not ship this
    p = np.exp(logits) / np.exp(logits).sum()
    return -np.log(p[y])

naive_ce(np.array([800.0, 799.0]), 0)        # nan — exp(800) overflows float64

The correct loss is finite here: class 0 is only e^1 ≈ 2.7 times more likely than class 1, so L = log(1 + e^(−1)) = 0.313. Two facts rescue it. Softmax is invariant to adding a shared constant to every logit, and log-sum-exp can be computed stably by subtracting the maximum first:

def stable_ce(logits, y):
    z = logits - logits.max()                # largest entry becomes 0
    lse = np.log(np.exp(z).sum())            # log-sum-exp of shifted logits
    return -z[y] + lse

stable_ce(np.array([800.0, 799.0]), 0)       # 0.3133 — finite, correct

That invariance is worth stating plainly: adding the same constant to every logit changes nothing in the loss — only differences between logits matter. A model can let its logits drift to enormous magnitudes during training with zero effect on the objective, which is exactly the regime where the naive formula dies and the shifted one does not.

PyTorch runs this fused computation, and the surrounding habits follow from it:

logits = model(xb)                   # final layer is Linear(d, K) — no softmax
loss = F.cross_entropy(logits, yb)   # yb: LongTensor of class indices, shape (N,)
pred = logits.argmax(dim=1)          # softmax would not change argmax

The beginner detour to know by name: putting a Softmax layer inside the model and calling CrossEntropyLoss anyway. Recent PyTorch accepts it — it log-softmaxes whatever you hand it — but you have already softmaxed, and the double squish compresses the logit differences that carry the learning signal. Gradients shrink, training crawls, and the loss curve still looks plausible, which is what makes the bug expensive. The rule to quote: the model emits logits; the loss owns the non-linearity. Softmax happens once — inside the loss during training, or after inference when you need probabilities — never both before and inside.

The same formula explains why this loss trains so well. Its gradient with respect to each logit is p_j − 1{j = y}: the true class’s logit is pushed up by (1 − p_y), every other logit is pulled down by its own probability. Confident and correct means near-zero gradient; wrong and confident means gradient of magnitude near 1. Training ends when the softmax saturates on the true class, not when the logits stop growing — logits can drift arbitrarily large at no cost to the objective, a property the next chapter’s failure tour exploits.

One more detail you will meet in real code: reduction. F.cross_entropy averages over the batch by default, so the logged loss is a per-example mean; reduction='none' returns the per-example vector when you must weight or mask examples yourself, and reduction='sum' fits gradient-accumulation schemes that defer the optimizer step. Mind the scale when you switch — gradients are proportional to the loss, so a sum-reduced loss on a large batch silently raises the effective learning rate.

Mismatched pairings almost never crash. They converge — to the wrong place, at the wrong speed, with a loss curve that looks busy. Three mismatches you should be able to name and repair on the spot.

MSE with a sigmoid output for classification is the classic. Its gradient with respect to the logit is (p − y)·σ′(z), and σ′ = p(1 − p) vanishes whenever the sigmoid saturates — including when it is confidently wrong, which is precisely when the example carries the most information. Cross-entropy cancels that derivative: its gradient with respect to the logit is simply p − y. Same representable predictions, completely different training dynamics, which is why the cross-entropy family is the default and MSE on 0/1 labels shows up mostly in legacy code. The double-softmax mistake from the previous chapter belongs in the same family: a mismatch that still trains, just slower, while the loss says nothing is wrong.

The second mismatch is direction, not speed: sigmoid-plus-BCE on mutually exclusive classes, or softmax-plus-CE on multi-label data. Softmax outputs must sum to 1, so feeding them to BCEWithLogitsLoss asks each one to be independently 0 or 1 — the objectives contradict each other, per-class losses stop being comparable, and the model receives conflicting gradients for every multi-tag example. Choose the pairing from the label structure: can two labels be true at once? Then sigmoid and BCEWithLogits. No? Then softmax and cross-entropy.

The third is imbalance, where the default loss is quietly wrong. With a 1% positive rate, an all-negative predictor scores 99% accuracy and its gradients for the positive class barely register. Two knobs rebalance the objective without touching the data:

# multiclass where class 1 is ~20x rarer than the others
crit = nn.CrossEntropyLoss(weight=torch.tensor([1.0, 20.0, 1.0]))

# binary with the same imbalance
crit = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([20.0]))

pos_weight multiplies the positive term of each example’s loss, so false negatives on the rare class dominate the gradient and the decision boundary shifts toward catching them. Be precise about the side effect: a weighted loss optimizes a reweighted distribution, so the probabilities that come out are no longer calibrated to the real class prior — recalibrate before presenting them as risks to anyone downstream. Focal loss, which down-weights easy well-classified examples, is the research-grade extension for extreme imbalance, but weighted cross-entropy is the answer that gets you the follow-up question. Whatever you weight, verify the calibration afterwards: bucket validation predictions by stated confidence and compare each bucket’s mean predicted probability against its actual frequency — systematic overshoot on the rare class is the signature of a reweighted loss living in a world where that class is commoner than it is. Isotonic or Platt calibration on a held-out set repairs the numbers without retraining, turning a weighted model’s outputs back into risks you can quote to a product manager.

Overconfidence rounds out the failure tour. With enough capacity, logits grow without bound, predicted confidence approaches 1.0, and a single mislabeled example near the boundary generates an enormous loss that distorts its whole batch. Two standard mitigations: label smoothing, nn.CrossEntropyLoss(label_smoothing=0.1), which spreads a little probability mass to wrong classes and measurably improves calibration for a fraction of a percent of accuracy; and regularization — weight decay plus early stopping — which keeps logit magnitudes in check for the reasons the previous chapter’s weight-growth caveat predicted.

Close with the drill that turns this topic into reflex. A product interviewer says: predict house price, churn probability, article topic, article tags. Your answers, in order: Linear(1) with Huber or MSE; Linear(1) with BCEWithLogitsLoss; Linear(K) with CrossEntropyLoss; Linear(K) with BCEWithLogitsLoss. Say them without pausing — then watch the same question reappear as the loss.backward() call in Backpropagation and the Training Loop.