🎓 A1Learner
🔍

PyTorch Logits vs Probabilities, the Classic Bug

Logits are the raw scores a model outputs before any squashing into probabilities. CrossEntropyLoss expects logits, and it applies its own log-softmax inside.

The buggy line

Softmax turns logits into probabilities. Probabilities are the right output for a person to read.

import torch

logits = torch.tensor([[2.0, 1.0, 0.0]])
probs = torch.softmax(logits, dim=1)
print(probs.round(decimals=4))

The same loss, two inputs

Cross-entropy scores the logits directly. Feeding it probabilities scores a different, softer input.

import torch

logits = torch.tensor([[2.0, 1.0, 0.0]])
target = torch.tensor([0])
loss_fn = torch.nn.CrossEntropyLoss()
print(f"{loss_fn(logits, target).item():.4f}")
probs = torch.softmax(logits, dim=1)
print(f"{loss_fn(probs, target).item():.4f}")

Softened twice

Cross-entropy applies log-softmax to its input. Probabilities fed in get squashed a second time and move closer together.

import torch

probs = torch.tensor([[0.6652, 0.2447, 0.0900]])
print(torch.softmax(probs, dim=1).round(decimals=4))

Example

import torch

torch.manual_seed(0)

# Three clusters of points, one per class.
centers = torch.tensor([[-3.0, 0.0], [3.0, -2.0], [0.0, 3.0]])
X = torch.cat([c + 0.8 * torch.randn(30, 2) for c in centers])
y = torch.cat([torch.full((30,), i) for i in range(3)])


def train(buggy):
    torch.manual_seed(1)
    model = torch.nn.Linear(2, 3)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.3)
    loss_fn = torch.nn.CrossEntropyLoss()
    trace = []
    for step in range(61):
        logits = model(X)
        if buggy:
            fed = torch.softmax(logits, dim=1)
        else:
            fed = logits
        loss = loss_fn(fed, y)
        if step % 20 == 0:
            trace.append(loss.item())
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    return trace


correct = train(False)
buggy = train(True)
for step, c, b in zip([0, 20, 40, 60], correct, buggy):
    print(f"Step {step}: correct {c:.4f}  buggy {b:.4f}")

Key points

  • Logits: the raw scores from the last layer, before softmax
  • Probability: a score squashed by softmax so the scores sum to 1
  • Use nn.CrossEntropyLoss with logits, never with probabilities
  • Softmax before cross-entropy softens the scores twice, and learning slows down
  • The fix for the bug is deleting the softmax line, nothing more
  • Raw scores go into cross-entropy, and the loss does its own squashing inside

In this example, three clusters of points form three classes. Two models start from the same seed and train on the same data. The correct model feeds logits straight into CrossEntropyLoss. The buggy model applies softmax first. At step 0 the buggy loss reads lower, at 1.2890 against 2.2989. By step 20 the correct loss has fallen to 0.0563, and the buggy loss still sits at 0.6173. At step 60 the gap remains, at 0.0284 against 0.5770. The double-softened scores give the loss far less to learn from. The buggy model crawls while the correct model finishes.

Note

nn.CrossEntropyLoss combines log-softmax and negative log-likelihood in one step

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

A model ends with softmax, and its output goes into nn.CrossEntropyLoss. What is the fix?