# Loss functions and the logits rule — PyTorch

Source: https://www.skillbyai.com/en/pytorch/pt-loss-functions

> Match each task to the right output layer and loss, and feed raw logits to CrossEntropyLoss and BCEWithLogitsLoss.

## Picking the right loss

For **multi-class classification** (exactly one correct class), use `nn.CrossEntropyLoss`. Its input is raw scores called **logits**, shape `(N, C)`, and its target is class indices with dtype `int64`, shape `(N,)`. It applies `log_softmax` internally and then the negative log-likelihood. Applying `softmax` yourself first is a classic bug: the model still trains, but gradients are weaker and results worse.

For **binary or multi-label** classification, use `nn.BCEWithLogitsLoss` with logits and **float** targets of the same shape. It combines sigmoid and binary cross-entropy in a numerically stable way, so prefer it to `nn.Sigmoid` followed by `nn.BCELoss`.

For **regression**, use `nn.MSELoss` (penalises large errors heavily) or `nn.L1Loss` / `nn.SmoothL1Loss` (more robust to outliers).

By default losses use `reduction="mean"`, averaging over the batch so the loss scale does not depend on batch size. Apply softmax or sigmoid only when you need probabilities for reporting or decisions, not before the loss.

## Cross-entropy and BCE with logits

Logits go straight into the loss. The targets differ: class indices (`int64`) for cross-entropy, floats of the same shape for BCE.

```python
import torch
from torch import nn

logits = torch.tensor([[2.0, 0.5, -1.0],
                       [0.1, 0.2, 3.0]])   # (N=2, C=3) raw scores
targets = torch.tensor([0, 2])            # class indices, int64

ce = nn.CrossEntropyLoss()
print(ce(logits, targets).item())

bce = nn.BCEWithLogitsLoss()
binary_logits = torch.tensor([1.5, -0.3])
binary_targets = torch.tensor([1.0, 0.0])  # floats, same shape as logits
print(bce(binary_logits, binary_targets).item())
```

## What cross-entropy computes, by hand

Per sample, cross-entropy is `log(sum(exp(logits))) - logit_of_true_class`. This plain-Python version reproduces the mean that `nn.CrossEntropyLoss` returns for the logits above.

```python
import math

logits = [[2.0, 0.5, -1.0], [0.1, 0.2, 3.0]]
targets = [0, 2]
losses = []
for row, t in zip(logits, targets):
    log_sum_exp = math.log(sum(math.exp(z) for z in row))
    losses.append(log_sum_exp - row[t])   # -log softmax(row)[t]
print([round(l, 4) for l in losses])
print(round(sum(losses) / len(losses), 4))
```

Output:

```
[0.2413, 0.1096]
0.1755
```

## Task to loss lookup

Multi-class: final `Linear(…, C)` + `CrossEntropyLoss`, `int64` targets. Binary or multi-label: final `Linear(…, 1 or K)` + `BCEWithLogitsLoss`, float targets. Regression: final `Linear(…, 1)` + `MSELoss` or `L1Loss`. No softmax or sigmoid layer at the end of the model in any of these.

**Quiz:** Your model outputs logits of shape `(N, 5)` and your labels are integers from 0 to 4. Which setup is correct?

- [ ] Apply softmax, then `nn.MSELoss` against the labels
- [ ] `nn.BCEWithLogitsLoss` with the integer labels
- [x] `nn.CrossEntropyLoss` on the raw logits with `int64` labels
- [ ] Apply softmax, then `nn.CrossEntropyLoss`

*Answer:* `nn.CrossEntropyLoss` on the raw logits with `int64` labels. CrossEntropyLoss expects raw logits and class-index targets; it applies log_softmax itself. Adding softmax first double-applies it and weakens gradients.