Logits, softmax and picking the right loss

Setup#

pip install torch

Outputs captured with torch 2.5.1 on CPU; these numbers are deterministic.

The bug, measured#

import torch
import torch.nn as nn

logits = torch.tensor([[ 3.0, -1.0,  0.5],
                       [ 0.2,  2.5, -0.3]])      # raw model outputs: two samples, three classes
targets = torch.tensor([0, 1])                   # class indices, dtype long

ce = nn.CrossEntropyLoss()
print(f"correct, raw logits in:       {ce(logits, targets).item():.4f}")
print(f"bug, softmax before the loss:  {ce(logits.softmax(dim=1), targets).item():.4f}")

z = torch.tensor([ 2.0, -1.0,  0.3])             # one output per sample: binary case
t = torch.tensor([ 1.0,  0.0,  1.0])             # binary targets, dtype float
print(f"BCEWithLogitsLoss on raw z:    {nn.BCEWithLogitsLoss()(z, t).item():.4f}")
print(f"BCELoss on sigmoid(z):         {nn.BCELoss()(z.sigmoid(), t).item():.4f}")

correct, raw logits in:       0.1225
bug, softmax before the loss:  0.6285
BCEWithLogitsLoss on raw z:    0.3315
BCELoss on sigmoid(z):         0.3315

Same predictions, same targets — and the softmax-first version reports a loss five times higher. Worse than the wrong number: probabilities are squashed between 0 and 1, so after the loss re-converts them, every prediction looks lukewarm. Confidence can never grow past a ceiling, gradients shrink, and training plateaus early while looking otherwise healthy.

Walkthrough#

CrossEntropyLoss = log-softmax + negative log likelihood, fused. The fusion is not only convenience — computing softmax and log separately invites numerical overflow for large logits, and the fused version sidesteps it. That is why the API insists on logits.

Targets are class indices, dtype long. Shape (N,) with values 0..C-1, not one-hot vectors. Handing it floats raises a dtype error — covered with its friends in dtype mismatch errors.

The binary twins. BCEWithLogitsLoss is sigmoid + binary cross-entropy fused; BCELoss expects you to sigmoid first. The output above shows they agree when used correctly. Prefer the fused one: same answer, better numerical safety, one less chance to forget.

At inference time, apply softmax when a human needs probabilities. For picking the winning class, skip it: logits.argmax(dim=1) — softmax never changes the ranking.

Choosing the loss#

Multi-label deserves the highlight: CrossEntropyLoss cannot express "this photo has a dog and a ball", because softmax forces classes to compete for one budget. Independent sigmoids, via BCEWithLogitsLoss, let each label stand alone.

Common mistakes#

softmax (or nn.Softmax) as the model's last layer, then CrossEntropyLoss. The measured bug above. Delete the softmax from the model.

sigmoid in the model, then BCEWithLogitsLoss. The binary edition of the same double-conversion. Either drop the sigmoid, or switch the loss to BCELoss — dropping the sigmoid is the better fix.

log_softmax into CrossEntropyLoss. The correct pairing for log-softmax outputs is NLLLoss. Feeding log-probabilities into a loss that log-softmaxes again shifts every loss value and weakens gradients, in the same quiet way.

Binary classification with two output neurons and mismatched machinery. Two logits with CrossEntropyLoss works; one logit with BCEWithLogitsLoss works. Mixing halves of the two setups produces shape errors, or worse, silent nonsense.

Try it yourself#

Extend the race from choosing an optimiser: add nn.Softmax(dim=1) as a final layer and retrain. Compare the loss after 200 steps against the clean model, and note that nothing ever crashed.

What to learn next#

• Writing a validation loop that reports the truth — measuring models without fooling yourself.
• Custom loss functions in PyTorch — when the built-in losses stop being enough.
• Loss functions — the conceptual tour behind these APIs.