A classifier can produce two scores for an image, but scores are not yet probabilities. To interpret them as the chances of “airplane” and “bird,” each output must lie between 0 and 1, and the two outputs must add up to 1. Softmax makes that conversion while remaining differentiable. section 7.2.3
Softmax exponentiates each score, then divides each result by the sum of all the exponentials. It preserves the scores’ ranking, though it changes their ratios. For scores [1, 2, 3], the result is approximately [0.0900, 0.2447, 0.6652]. section 7.2.4
# In[10]:
softmax = nn.Softmax(dim=1)
x = torch.tensor([[1.0, 2.0, 3.0],
[3.0, 2.0, 1.0]])
softmax(x)
# Out[10]:
tensor([[0.0900, 0.2447, 0.6652],
[0.6652, 0.2447, 0.0900]])
Two rows of scores show why nn.Softmax(dim=1) matters: it normalizes each row separately, as it would for a batch of predictions. section 7.2.4
For image classification, a final nn.Softmax can turn a network’s two output scores into class probabilities. The class with the larger probability is its prediction. For training, the related nn.LogSoftmax can be paired with nn.NLLLoss; the softmax-and-negative-log-likelihood combination is cross-entropy loss. section 7.2.4 section 7.5
Softmax also has a different job in self-attention. Applied to query–key dot-product scores, it produces attention weights that sum to 1. Those weights determine how strongly value vectors contribute to a weighted sum. section 9.5.1
