the last three articles walked the training loop backwards, .backward() and what it fills in, the four gradients computed by hand, then optimizer.step() turning those gradients into an actual update. all of that starts at one number, the loss, and i kept saying "and then the loss comes from somewhere" and moving on. so this one is about where it comes from, which is a much shorter piece of arithmetic than you would expect, and then about the one mistake in this area that every beginner is warned about and almost nobody has actually measured.
the warning is "never apply softmax before CrossEntropyLoss", and you will find it in every tutorial, every course and most stack overflow answers on the subject, almost always with the words "training will fail quietly". i wanted a number for "quietly", so i ran it. it does not fail. it trains, it reaches roughly the same accuracy, and it lies to you the whole way in a specific and measurable manner, which turns out to be a lot more interesting than failing would have been.
The short answer
if the answer your model produces is a number, use MSELoss. if the answer is a category, use CrossEntropyLoss. binary yes or no, BCEWithLogitsLoss. that is the whole decision for the vast majority of problems, and neither of the two main ones has anything hidden inside it:
MSE = ((pred - target) ** 2).mean()
CrossEntropy = -log( softmax(logits)[correct_class] )
the loss function is the only place in the entire loop where you tell the model what "better" means. the architecture, the optimizer, the data, all of it is machinery in service of that one definition.
MSE, which is exactly what it says
take the difference, square it so negative errors do not cancel positive ones, average over the batch.
import torch
import torch.nn as nn
pred = torch.tensor([2.5, 0.5, 3.0])
target = torch.tensor([3.0, 0.0, 3.0])
print(nn.MSELoss()(pred, target)) # tensor(0.1667)
print(((pred - target) ** 2).mean()) # tensor(0.1667)
by hand the errors are (-0.5, 0.5, 0.0), squared they are (0.25, 0.25, 0.0), and 0.5 / 3 is 0.1667. the nn.MSELoss() object is a wrapper around that one line and nothing else.
the squaring is the part with consequences. shift all your predictions by 1.0 and the loss is 1.0, shift them by 2.0 and it is 4.0, so an error twice as large costs four times as much and MSE will always spend most of its effort on your worst outliers. that is either what you want or the reason your model is being dragged around by three bad rows in the dataset.
the MSE trap that costs an afternoon
this one is not in the docs where you would look for it. give MSE a prediction of shape (4, 1) and a target of shape (4,), which happens the moment you use nn.Linear(n, 1) and your labels came out of a dataframe column:
pr = torch.tensor([[1.0], [2.0], [3.0], [4.0]]) # (4, 1)
tg = torch.tensor([1.0, 2.0, 3.0, 4.0]) # (4,)
print(nn.MSELoss()(pr, tg)) # tensor(2.5000)
the prediction is exactly right on every row and the loss is 2.5. broadcasting expanded (4,1) against (4,) into a (4,4) grid and compared every prediction against every target, so twelve of the sixteen comparisons are between rows that have nothing to do with each other. pytorch does warn here, and the warning is easy to scroll past because your training loop still prints a falling loss:
UserWarning: Using a target size (torch.Size([4])) that is different to the input size
(torch.Size([4, 1])). This will likely lead to incorrect results due to broadcasting.
Please ensure they have the same size.
fix is target.unsqueeze(1) or pred.squeeze(1), and then the loss is 0.0 like it should be. if the (4,4) part is surprising, it is the same three broadcasting rules from the broadcasting article, just showing up somewhere you were not looking for them.
Cross-entropy, by hand
classification is a different shape of problem. the model outputs one raw score per class, called logits, and the target is an integer saying which class was right. cross-entropy turns the scores into probabilities with softmax, then takes the negative log of the probability sitting on the correct class. more probability on the right answer, lower loss.
logits = torch.tensor([
[2.0, 1.0, 0.1, -1.0], # sample 0, correct class 0
[0.5, 2.5, 0.3, 0.2], # sample 1, correct class 1
[0.1, 0.2, 3.0, 0.1], # sample 2, correct class 2
])
labels = torch.tensor([0, 1, 2])
print(nn.CrossEntropyLoss()(logits, labels)) # tensor(0.3015)
now the same thing by hand for sample 0:
l = torch.tensor([2.0, 1.0, 0.1, -1.0])
p = torch.softmax(l, dim=0)
print(p) # tensor([0.6381, 0.2347, 0.0954, 0.0318])
print(p.sum()) # tensor(1.0000)
print(-torch.log(p[0])) # tensor(0.4493)
softmax exponentiates every score and divides by the total, so the four numbers come out positive and sum to one. class 0 got 63.8 percent of the probability and -log(0.6381) is 0.4493. the other two samples give 0.2974 and 0.1577, and the mean of the three is 0.3015, which is what the loss function returned.
one habit worth building right here. whenever you print a softmax, print its sum on the next line. it costs four characters and it is the only cheap check that exists, because a vector that does not add up to 1 is not a probability distribution and something upstream of it is wrong, usually a softmax taken over the wrong dimension. i have run into more than one published table of "probabilities" that quietly fails that test, and the sum is what catches it every time.
why the negative log, and not something simpler
because of what it does to the gradient. compute the gradient of cross-entropy with respect to the logits and you get one of the cleanest results in the whole subject:
lg = torch.tensor([[2.0, 1.0, 0.1, -1.0]], requires_grad=True)
nn.CrossEntropyLoss()(lg, torch.tensor([0])).backward()
print(lg.grad) # tensor([[-0.3619, 0.2347, 0.0954, 0.0318]])
that is p - y, the predicted probability vector minus the one-hot target, exactly, to every decimal place. 0.6381 - 1 = -0.3619 on the correct class and the raw probability on each of the others. the softmax and the log cancel each other's derivatives and what reaches your network is just "how far off was each probability". if you followed the backpropagation article, this is the number that starts the whole chain, and every gradient in the model is that vector pushed backwards through the layers.
it also means the size of the gradient is bounded by 1 per class and it goes to zero exactly when the prediction is right. no tuning, no scaling, it just behaves.
The softmax bug, actually measured
here is the thing everyone tells you. CrossEntropyLoss applies softmax internally, so if you also apply softmax in your model's forward, you softmax twice. everyone repeats it and i have repeated it, and almost nobody says what it actually costs, so that is what i went and measured.
so, first, it does not fail:
correct = nn.CrossEntropyLoss()(logits, labels)
bugged = nn.CrossEntropyLoss()(torch.softmax(logits, dim=1), labels)
print(correct) # tensor(0.3015)
print(bugged) # tensor(0.9388)
no error, no warning, two perfectly ordinary numbers. and there is no way for pytorch to catch this, because a probability vector is a valid float tensor and "valid logits" is not a thing you can test for.
what happened is that softmax ran on top of softmax. sample 0's logits are [2.0, 1.0, 0.1, -1.0], spread over three units. after one softmax they are [0.6381, 0.2347, 0.0954, 0.0318], spread over about 0.6. after the second one they are [0.3578, 0.2391, 0.2080, 0.1951], spread over 0.16. every pass through softmax squeezes the numbers closer together, and the model's confidence is what gets squeezed out.
what that costs, in one number
i pushed it to the extreme, a model that is maximally confident and correct, and the same model maximally confident and wrong, and asked both losses what they think.
confident + right confident + wrong
logits in 0.0000 40.0000
softmax in 1.4612 2.4612
that is with 10 classes. a correctly wired cross-entropy has a floor at zero and no ceiling at all, so being badly wrong is expensive without limit and the loss carries real information about how wrong you are. once you softmax first, the entire range of model behaviour from perfect to catastrophic gets compressed into a band exactly 1.0 wide. i checked this for 3, 4, 10 and 100 classes and the band is 1.0000 every single time, it only slides upward as the class count grows.
so the loss stops being a measurement. and the gradient goes with it:
gradient on the true-class logit of a confidently wrong sample, 10 classes
logits in -1.000e+00
softmax in -4.871e-18
seventeen orders of magnitude. the example your model is most badly wrong about, the one it most needs to learn from, produces a gradient of effectively zero. that is not a rounding artifact, it is the second softmax having saturated, and the derivative of a saturated softmax is nothing.
and yet it trains
this is the part i did not expect and it is the reason i think the standard warning is told wrong. ten classes, twenty dimensions, same network, same seed, same optimizer, the only difference being one torch.softmax in the forward pass:
epoch 1 10 50 200 1000 4000 final test acc
------------------------------------------------------------------------------------
logits in 20.3% 75.9% 85.6% 85.9% 85.0% 82.5% 82.5%
softmax in 12.6% 24.7% 62.4% 85.6% 86.4% 85.4% 85.4%
the bugged run is far slower, 24.7 percent against 75.9 at epoch 10, and it needs roughly four times as long to get anywhere. but it gets there. it finishes ahead, in fact, because the correct run had started overfitting by epoch 4000 and the crushed gradients acted as a brake. i am not claiming the bug is good, on a real problem with a real budget being four times slower is the whole ball game, and i would not want to defend "my regularizer is a bug" to anyone. the point is that "training will fail quietly" sets you up to look for a failure, and there is no failure to find.
what you get instead is worse in a more annoying way. your loss is now uninterpretable. that final softmax in run reports a loss of 1.53. with 10 classes and correct wiring, 1.53 would mean a mediocre model, since a model guessing at random sits at log(10) = 2.3026 and 1.53 is two thirds of the way to useless. with the bug, 1.53 is 0.07 above the best value that run can physically produce. same number, opposite meanings, and nothing on your screen says which one you are looking at.
and accuracy will not save you either, because softmax is monotone, so argmax does not move and your accuracy metric is completely unaffected by the bug. i checked it over a thousand random rows, the predicted class is identical in every one. the only two things that change are the number you watch and the speed you learn, which is exactly the pair you would least like to have quietly corrupted.
how to spot it in thirty seconds
look at the last line of your forward. if it is a softmax, a sigmoid, or anything else that squashes into a range, and your loss is CrossEntropyLoss or BCEWithLogitsLoss, that is the bug. those two losses do the squashing themselves, in a numerically stable way, which is a second reason to let them.
the other check is the floor. run a batch you know the model gets right and see whether the loss can reach something near zero. if the lowest loss you ever see is suspiciously far from zero and suspiciously close to a constant, you are looking at a band, not a measurement.
# wrong
class Net(nn.Module):
def forward(self, x):
return torch.softmax(self.fc(x), dim=1) # then CrossEntropyLoss -> bug
# right
class Net(nn.Module):
def forward(self, x):
return self.fc(x) # raw logits, that is all
# and at inference time, when you want probabilities to show a user:
with torch.no_grad():
probs = torch.softmax(model(x), dim=1)
softmax at prediction time is normal and correct. softmax before the loss is the bug. same function, and the only thing that differs is which side of the loss it sits on.
Why not just use MSE for classification
people ask this, and the answer is the same gradient story. take a model that is confidently wrong on a 3-class problem, logits [-6, 0, 0] with the true class being 0, and look at the gradient on the true-class logit under each loss:
cross-entropy -0.998762
MSE on softmax outputs -0.00123478
809 times smaller. MSE on a squashed output inherits the squashing, so the more wrong the model is the less it learns, which is precisely backwards. cross-entropy's p - y is at its largest exactly when the model is at its most wrong. that is the entire reason cross-entropy is the classification loss and not just a convention.
The error messages, verified
these are the literal strings from torch 2.14, not from memory, because the ones you find on stack overflow are frequently from 2019.
| what you did | what you get |
|---|---|
float labels into CrossEntropyLoss
|
RuntimeError: expected target dtype to be Long or Byte, but got Float |
a label larger than num_classes - 1
|
IndexError: Target 9 is out of bounds. |
| batch sizes do not match | ValueError: Expected input batch_size (3) to match target batch_size (2). |
raw logits into BCELoss
|
RuntimeError: all elements of input should be between 0 and 1 |
softmax into CrossEntropyLoss
|
nothing at all, which is the whole article |
that last row is why this piece exists. four of the five mistakes stop your program. the fifth one promotes itself to a training run.
worth noting that the BCELoss row is the friendlier cousin of the softmax bug. BCELoss genuinely requires probabilities and BCEWithLogitsLoss requires raw scores, and if you mix those up in the direction of feeding logits to BCELoss you at least get told, because a logit of 1.5 is outside [0, 1] and that is checkable. the other direction, sigmoid then BCEWithLogitsLoss, is silent in exactly the way described above.
Reduction, briefly
by default every loss averages over the batch, and you almost always want that, because it makes your learning rate independent of batch size.
pred = torch.tensor([1.0, 2.0, 5.0])
tgt = torch.tensor([1.0, 3.0, 3.0])
nn.MSELoss(reduction='mean')(pred, tgt) # tensor(1.6667)
nn.MSELoss(reduction='sum')(pred, tgt) # tensor(5.)
nn.MSELoss(reduction='none')(pred, tgt) # tensor([0., 1., 4.])
sum scales your effective learning rate with batch size, which is a real bug source when someone changes the batch size and training suddenly diverges. none is genuinely useful, it gives you the per-sample loss, which is how you weight samples differently or go find which rows your model hates.
The rule
regression, MSELoss. multi-class, CrossEntropyLoss on raw logits. binary or multi-label, BCEWithLogitsLoss on raw logits. never squash before a loss whose name contains "Logits", and CrossEntropyLoss counts even though its name does not say so. MSE is ((pred - target) ** 2).mean() and cross-entropy is -log(softmax(logits)[label]), and if you can write both of those from memory you understand loss functions well enough to move on.
Try it before you close the tab
take the ten-class setup, or any classifier you already have, and add torch.softmax(x, dim=1) to the end of the forward pass. do not change anything else. then train it and look only at the loss curve and tell me you could spot which one is broken. that experiment took me about four minutes and it is the reason i no longer trust a falling loss on its own.
then print the smallest loss your run ever reaches and compare it to log(num_classes). a healthy run gets near zero on data it has memorised. a run stuck at a constant well above zero is telling you something, and this bug is one of the things it could be telling you.
what is the loss value you would consider "good" on your current problem, and do you know why that number and not a different one? i asked myself that while writing this chapter and my honest answer was that i had been reading loss curves by their shape for a long time and had never once checked what the floor should be.
This is one chapter's worth of an idea from my book, PyTorch From Ground Up, which builds everything from tensors upward so nothing stays vague. If it helped: 8 chapters are free, no email required, there's a free one-page tensor cheat-sheet here, every example runs in the companion notebooks on GitHub, and the full book is on Leanpub or in paperback and Kindle on Amazon.
More in this series
How Training Actually Works, the part where the training loop stops being magic:
- PyTorch Autograd Explained: What
.backward()Actually Does - Backpropagation by Hand: Two Layers, a Pen, and Then Autograd Agrees
- What
optimizer.step()Actually Does: SGD, Momentum and Adam by Hand
The shape mechanics underneath all of it, worth having solid first:
- Reshape vs View in PyTorch
- PyTorch Broadcasting Explained
- What Does
unsqueezeDo in PyTorch? - What Does
keepdimDo in PyTorch? - PyTorch
permutevstranspose
Coming next in How Training Actually Works: building the same network twice, once from raw tensors and once from nn.Module, and what the module system is actually keeping track of for you.
Top comments (0)