PyTorch Softmax: dim, log_softmax and CrossEntropyLoss

torch.softmax turns raw model scores into probabilities that add up to 1:

probs = torch.softmax(scores, dim=1)     # dim is REQUIRED

Two things go wrong with it constantly. Picking the wrong dim normalises across the batch instead of across the classes, and applying softmax before CrossEntropyLoss applies it twice.

Both are demonstrated below, on Python 3.12.5, PyTorch 2.14.0+cpu.

Using softmax in PyTorch

Every value becomes positive, and each group sums to exactly 1:

import torch

scores = torch.tensor([[1.0, 2.0, 3.0],
                       [1.0, 1.0, 1.0]])

probs = torch.softmax(scores, dim=1)

print(probs)
print()
print("each row sums to:", probs.sum(dim=1))
print("equal scores give equal probabilities, as the second row shows")

Output:

tensor([[0.0900, 0.2447, 0.6652],
        [0.3333, 0.3333, 0.3333]])

each row sums to: tensor([1., 1.])
equal scores give equal probabilities, as the second row shows

The gaps get exaggerated. Scores of 1, 2 and 3 become roughly 0.09, 0.24 and 0.67, because softmax exponentiates before it normalises.

Equal scores give equal probabilities, which is the sanity check in the second row of that output.

torch.softmax vs F.softmax vs nn.Softmax

Three spellings, one computation. The difference is where you use them:

import torch
import torch.nn as nn
import torch.nn.functional as F

x = torch.tensor([[1.0, 2.0, 3.0]])

a = torch.softmax(x, dim=1)        # the tensor function
b = F.softmax(x, dim=1)            # the functional form
c = nn.Softmax(dim=1)(x)           # the module form

print("torch.softmax :", a)
print("F.softmax     :", b)
print("nn.Softmax    :", c)

print()
print("identical results:", torch.allclose(a, b) and torch.allclose(b, c))
print("but torch.softmax IS F.softmax?", torch.softmax is F.softmax)
print("  (different objects, same kernel underneath)")

Output:

torch.softmax : tensor([[0.0900, 0.2447, 0.6652]])
F.softmax     : tensor([[0.0900, 0.2447, 0.6652]])
nn.Softmax    : tensor([[0.0900, 0.2447, 0.6652]])

identical results: True
but torch.softmax IS F.softmax? False
  (different objects, same kernel underneath)

Form	Typical use
torch.softmax(x, dim)	Anywhere, shortest to type
F.softmax(x, dim)	Inside a forward method
nn.Softmax(dim)	As a layer in nn.Sequential
x.softmax(dim)	The tensor method, same thing again

They produce identical numbers. nn.Softmax is a module wrapping the function, so it’s the one that can sit in a model definition alongside other layers.

Pick whichever reads best where you are. If softmax has to appear in a Sequential, you need the module form; otherwise the function is simpler.

What the softmax dim argument does

This is the bug that produces a model which trains and never learns:

import torch

# a batch of 2 samples with 3 class scores each
batch = torch.tensor([[1.0, 2.0, 3.0],
                      [1.0, 1.0, 1.0]])

print("dim=1, across each row (CORRECT for a batch):")
print(torch.softmax(batch, dim=1))
print("row sums:", torch.softmax(batch, dim=1).sum(dim=1))

print()
print("dim=0, down each column (almost always WRONG):")
print(torch.softmax(batch, dim=0))
print("column sums:", torch.softmax(batch, dim=0).sum(dim=0))

print()
print("dim=-1 means the last dimension, which is the safe habit")
print("same as dim=1 here:", torch.allclose(torch.softmax(batch, dim=-1),
                                            torch.softmax(batch, dim=1)))

Output:

dim=1, across each row (CORRECT for a batch):
tensor([[0.0900, 0.2447, 0.6652],
        [0.3333, 0.3333, 0.3333]])
row sums: tensor([1., 1.])

dim=0, down each column (almost always WRONG):
tensor([[0.5000, 0.7311, 0.8808],
        [0.5000, 0.2689, 0.1192]])
column sums: tensor([1.0000, 1.0000, 1.0000])

dim=-1 means the last dimension, which is the safe habit
same as dim=1 here: True

For a batch shaped (batch_size, num_classes) you want dim=1, so each sample gets a probability distribution over its classes.

dim=0 normalises down the batch, which compares sample 1’s cat score against sample 2’s cat score. That is meaningless.

dim=-1 is the safer habit. It always means the last dimension, so it keeps working when a sequence dimension appears and the shape becomes (batch, seq, classes). The same reasoning applies to nn.Linear, which also acts on the last dimension.

Why does PyTorch softmax require dim?

Leaving it out is an error rather than a guess:

import torch

x = torch.tensor([[1.0, 2.0, 3.0]])

try:
    torch.softmax(x)
except TypeError as err:
    print("torch.softmax(x) ->", type(err).__name__ + ":", err)

print()
print("dim is required, and that is deliberate")
print("guessing it silently would produce wrong numbers instead of an error")

Output:

torch.softmax(x) -> TypeError: softmax() missing 1 required positional arguments: "dim"

dim is required, and that is deliberate
guessing it silently would produce wrong numbers instead of an error

Older PyTorch versions guessed a dimension and warned about it. Requiring the argument was the right call, because a silently wrong axis is far worse than a stopped script.

Softmax and numerical stability

Computing the formula yourself overflows on large numbers. The built-in does not:

import torch

# large logits overflow if you compute exp() yourself
big = torch.tensor([1000.0, 1001.0, 1002.0])

print("torch.exp(big)     :", torch.exp(big), " <- overflowed")
print("torch.softmax(big) :", torch.softmax(big, dim=0), " <- fine")

print()
print("softmax subtracts the maximum first, so the exponent never blows up")
print("the answer is the same, because that shift cancels out")

Output:

torch.exp(big)     : tensor([inf, inf, inf])  <- overflowed
torch.softmax(big) : tensor([0.0900, 0.2447, 0.6652])  <- fine

softmax subtracts the maximum first, so the exponent never blows up
the answer is the same, because that shift cancels out

PyTorch subtracts the largest value before exponentiating. That shift cancels out in the division, so the answer is unchanged and nothing overflows.

It’s a good reason never to hand-roll softmax from the formula, even though the formula is only three lines.

log_softmax and why it exists

Taking the log of a softmax separately is a second chance to lose precision:

import torch

x = torch.tensor([-1000.0, -1001.0])

print("log(softmax(x)) :", torch.log(torch.softmax(x, dim=0)))
print("log_softmax(x)  :", torch.log_softmax(x, dim=0))

print()
print("they agree here, but log_softmax does it in one stable step")
print("taking log of a probability that underflowed to 0 gives -inf")
print("which is why every loss function uses log_softmax internally")

Output:

log(softmax(x)) : tensor([-0.3133, -1.3133])
log_softmax(x)  : tensor([-0.3133, -1.3133])

they agree here, but log_softmax does it in one stable step
taking log of a probability that underflowed to 0 gives -inf
which is why every loss function uses log_softmax internally

If a probability underflows to zero, its log is -inf, and the gradient goes with it. log_softmax computes both steps together and never produces that zero.

Any time you need log-probabilities, which is most of classification, reach for log_softmax rather than composing the two.

Do not apply softmax before CrossEntropyLoss

This is the single most common PyTorch mistake, and nothing warns you about it:

import torch
import torch.nn as nn

logits = torch.tensor([[2.0, 1.0, 0.1]])
target = torch.tensor([0])

correct = nn.CrossEntropyLoss()(logits, target)
wrong = nn.CrossEntropyLoss()(torch.softmax(logits, dim=1), target)

print(f"CrossEntropyLoss(logits)          : {correct.item():.6f}   correct")
print(f"CrossEntropyLoss(softmax(logits)) : {wrong.item():.6f}   WRONG")

print()
# CrossEntropyLoss is NLLLoss plus log_softmax
manual = nn.NLLLoss()(torch.log_softmax(logits, dim=1), target)
print(f"NLLLoss(log_softmax(logits))      : {manual.item():.6f}   same as correct")

print()
print("CrossEntropyLoss applies log_softmax itself, so feeding it")
print("probabilities applies softmax TWICE. It trains, badly, with no error.")

Output:

CrossEntropyLoss(logits)          : 0.417030   correct
CrossEntropyLoss(softmax(logits)) : 0.802121   WRONG

NLLLoss(log_softmax(logits))      : 0.417030   same as correct

CrossEntropyLoss applies log_softmax itself, so feeding it
probabilities applies softmax TWICE. It trains, badly, with no error.

nn.CrossEntropyLoss applies log_softmax internally. Feeding it probabilities squashes an already-squashed distribution, flattening your gradients.

The model still trains. It just learns slowly and plateaus early, which is why this one goes unnoticed for so long.

• Pass raw logits to nn.CrossEntropyLoss.
• Pass log-probabilities to nn.NLLLoss, from log_softmax.
• Apply softmax only at inference, when you want readable probabilities.
• For a single-label prediction, argmax on the logits gives the same answer with no softmax at all.

The last point is worth remembering: softmax preserves order, so it never changes which class wins. See cross entropy loss for the full picture.

Turning PyTorch logits into probabilities

At inference time softmax is what makes the output readable:

import torch
import matplotlib.pyplot as plt
import numpy as np

logits = torch.tensor([2.0, 1.0, 0.5, 0.1])
labels = ["cat", "dog", "bird", "fish"]
probs = torch.softmax(logits, dim=0)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4.2))

ax1.bar(labels, logits.numpy(), color="#0b6bcb")
ax1.set_title("raw logits: any value, no meaning on their own")
ax1.axhline(0, color="#555", lw=.8)

ax2.bar(labels, probs.numpy(), color="#0a7d32")
ax2.set_title(f"after softmax: sums to {probs.sum():.0f}")
ax2.set_ylim(0, 1)

for ax in (ax1, ax2):
    ax.set_axisbelow(True); ax.grid(axis="y", alpha=.3)
plt.tight_layout(); plt.show()

for label, p in zip(labels, probs):
    print(f"  {label:<5} {p:.4f}")

Output:

  cat   0.5745
  dog   0.2114
  bird  0.1282
  fish  0.0859

Raw logits mean nothing on their own. A score of 2.0 is only meaningful next to the other scores in the same row.

Softmax puts them on a scale anyone can read, which is what you want for a confidence figure in an application.