Output heads and loss functions

The final layer of a network turns features into the thing you want to predict. Its shape and activation depend on the task, and so does the loss that trains it. Pick the wrong pair and the model can still train without an error. It just answers the wrong question.

The big idea

Softmax turns a list of scores into probabilities that are non-negative and sum to 1. Cross-entropy then measures how much probability went to the correct class.

Walking through softmax and cross-entropy

Take two classes with logits z = [2.0, 0.5]. Softmax exponentiates each logit, then divides by the total:

e^2.0 = 7.389,  e^0.5 = 1.649,  total = 9.038
p = [7.389 / 9.038, 1.649 / 9.038] ≈ [0.82, 0.18]

If the true class is 0, the loss is −log(0.82) ≈ 0.20. A network that outputs [0, 0] gives [0.5, 0.5], and its loss is ln 2 ≈ 0.69. That is the loss of random guessing for two classes, a useful baseline to check at the start of training.

The code below does this. It subtracts the maximum before exponentiating, which keeps large logits from overflowing. It runs as written.

import numpy as np

def softmax(z, T=1.0):
    z = z / T
    z = z - z.max()          # subtract the max so exp() cannot overflow
    e = np.exp(z)
    return e / e.sum()

z = np.array([2.0, 0.5])
p = softmax(z)
print(p.round(3))
print("CE if label is class 0:", round(-np.log(p[0]), 3))
print("uniform CE:", round(-np.log(softmax(np.zeros(2))[0]), 3))

Output:

[0.818 0.182]
CE if label is class 0: 0.201
uniform CE: 0.693

Temperature

Softmax can be sharpened or softened with a temperature T: softmax(z / T). A value below 1 makes the distribution more peaked, and a value above 1 flattens it. Training normally uses T = 1. Sampling and distillation use other values.

Common mistakes

• Using MSE on class labels. Squared error on 0 and 1 targets is a poor probabilistic objective, and the outputs can leave the range from 0 to 1.
• Using softmax for multi-label tasks. Softmax forces one winner. Independent sigmoids let several labels be true.
• Applying softmax and then taking the log separately. This can lose precision. Use the fused loss, such as cross_entropy in PyTorch, which takes raw logits.

Check yourself

1. Three classes, exactly one correct. Which head and loss? Answer: Three logits, softmax, and cross-entropy.
2. A photo can carry both the tags "beach" and "sunset". Softmax or sigmoids? Answer: Independent sigmoids, one per tag, because the tags are not mutually exclusive.

Key takeaways

• The output head must match the task, and the loss must match the head.
• Softmax plus cross-entropy is the standard pair for one-of-K classification.
• Use independent sigmoids when labels can overlap.
• Compare the loss at initialization with the random-guess value, ln K, as a sanity check.

Further learning

Probability and information theory explain why cross-entropy is the right loss: probability, statistics and information.

Next, count the parameters and see what depth and width buy: parameter counting, width and depth.