Unit 04.02: Regression and classification heads

The same network body serves regression, binary classification and multiclass. Only the last layer and the loss change.

Three heads, three output shapes

For regression the head emits one raw number with no activation — the answer is unbounded.

For binary classification it emits one logit. Apply the sigmoid at prediction time, not inside the model, because BCEWithLogitsLoss expects raw logits and is numerically safer than sigmoid followed by BCELoss.

For multiclass it emits one logit *per class*, and CrossEntropyLoss applies the softmax internally.

import torch
from torch import nn

torch.manual_seed(0)
features = torch.randn(6, 4)
body = nn.Sequential(nn.Linear(4, 16), nn.ReLU())
hidden = body(features)

# Regression: one raw number, no activation on the output.
regression_head = nn.Linear(16, 1)
print("regression out:", tuple(regression_head(hidden).shape), "unbounded")

# Binary: one logit. Apply sigmoid at prediction time, not inside the model,
# because BCEWithLogitsLoss expects raw logits and is numerically safer.
binary_head = nn.Linear(16, 1)
logits = binary_head(hidden)
print("binary logits :", tuple(logits.shape), "-> probs", [round(v, 3) for v in torch.sigmoid(logits).flatten()[:3].tolist()])

# Multiclass: one logit PER CLASS. Softmax is applied by CrossEntropyLoss.
multiclass_head = nn.Linear(16, 5)
scores = multiclass_head(hidden)
probs = torch.softmax(scores, dim=1)
print("multiclass    :", tuple(scores.shape), "row sums to", round(probs[0].sum().item(), 6))
print("predicted class:", probs.argmax(dim=1).tolist())

The multiclass probabilities sum to 1.0 across each row, which is what softmax guarantees, and argmax turns them into a predicted class.

Notice that the model itself never applies sigmoid or softmax. That is deliberate: the loss functions do it, more stably.

The mistake this prevents

Putting a softmax on the final layer and then using CrossEntropyLoss. The softmax is applied twice, gradients shrink, and training crawls for no visible reason.

Takeaway

Pick the head from the target type, leave the final activation to the loss function.

Mark complete