Softmax

How This Lesson Fits Module 6.1

Softmax is the multiclass counterpart to sigmoid. It belongs near the output layer, where raw class logits become an interpretable distribution over mutually exclusive classes.

Learning Objectives

By the end of this lesson, students should be able to:

Define softmax as a normalization over class logits.
Explain why softmax outputs sum to 1.
Distinguish multiclass softmax from multilabel sigmoid.
Use CrossEntropyLoss correctly with raw logits.
Convert logits to probabilities for inference.
Recognize numerical-stability and confidence-calibration issues.

Definition

Softmax exponentiates each class logit and divides by the sum across classes, producing nonnegative probabilities that add to 1.

A Probability Distribution Over Classes

For exclusive classes, the model should assign probability mass across competing options. Raising one class probability lowers the others. Softmax does this by comparing logits relative to each other, not in isolation. During PyTorch training, nn.CrossEntropyLoss expects raw logits and internally applies a stable log-softmax plus negative log-likelihood.

Problem type	Final scores	Activation/loss pattern
Binary	One logit	`BCEWithLogitsLoss`; sigmoid for inference
Multiclass	One logit per class	`CrossEntropyLoss`; softmax for inference
Multilabel	One logit per label	`BCEWithLogitsLoss`; sigmoid per label
Ranking	Scores per candidate	Softmax or ranking loss depending on setup
Calibration	Probabilities	Temperature scaling may help

PyTorch Practice

Apply softmax for reporting probabilities, but feed raw logits to CrossEntropyLoss during training.

import torch
from torch import nn

logits = torch.tensor([[2.0, 0.5, -1.0], [0.1, 1.2, 0.3]])
target = torch.tensor([0, 1])

loss_fn = nn.CrossEntropyLoss()
loss = loss_fn(logits, target)

probs = torch.softmax(logits, dim=1)
pred = probs.argmax(dim=1)
print(probs)
print(pred, loss.item())

Softmax vs Sigmoid

Classes compete
Outputs sum to 1
Use for single-label multiclass

Sigmoid

Labels independent
Each output is 0–1
Use for binary or multilabel

Shared caution

Do not duplicate stable loss internals
Use logits for training
Calibrate probabilities if decisions are high stakes

Strengths and Tradeoffs

Useful because

Produces an interpretable distribution for exclusive classes.
Pairs cleanly with cross-entropy training.
Argmax gives a simple predicted class.

Watch for

Can be overconfident even when wrong.
Not appropriate when multiple labels can be true together.
Naive exponentiation can overflow without stable implementations.

How It Flows

1. Logits

Output layer emits one raw score per class.

2. Normalize

Softmax converts relative scores to probabilities.

3. Select

Argmax chooses the largest probability for top-1 prediction.

4. Evaluate

Metrics compare predictions and confidence against targets.

Common Misconception

Do not use softmax for multilabel classification where several labels may be true at once. Softmax forces competition; multilabel tasks need independent sigmoid outputs.

Knowledge Check

Short Answer: What do softmax outputs sum to? Answer: 1.
True/False: Softmax is suitable for mutually exclusive classes. Answer: True.
Multiple Choice: PyTorch CrossEntropyLoss expects: (a) logits, (b) softmax probabilities, (c) strings. Answer: (a).
Short Answer: Which dimension is usually classes for [N,C]? Answer: Dimension 1.
True/False: Softmax labels are independent. Answer: False; they compete.
Short Answer: How get predicted class from probabilities? Answer: Argmax over class dimension.
Multiple Choice: Multilabel output usually uses: (a) sigmoid, (b) one shared softmax, (c) no loss. Answer: (a).
Short Answer: Why use stable loss implementations? Answer: To avoid numerical overflow/underflow and improve gradients.
True/False: Softmax probabilities can be poorly calibrated. Answer: True.
Short Answer: What is a logit? Answer: A raw unnormalized class score.

Key Takeaways

Softmax converts class logits into a distribution over exclusive classes.
Use raw logits with CrossEntropyLoss during training.
Use sigmoid, not softmax, when labels are independent.
Next, Loss Functions explains the objective that drives learning.

Trainer’s Guide

Hands-on idea: Ask students to compute softmax for three logits by hand after subtracting the max logit for stability.

Discussion prompt: Why can a model be accurate but still poorly calibrated?

Recap: Softmax is the interpretation layer for multiclass logits, while cross-entropy handles stable training from raw scores. Continue with Loss Functions.