Idée
Une fonction qui prend dix nombres bruts (pouvant être positifs ou négatifs, grands ou petits) et les transforme en dix nombres positifs qui somment à 1 — donc en une distribution de probabilités sur les dix chiffres possibles.
Outil
Comme l'exponentielle suivie d'une normalisation : $\sigma(z)_i = e^{z_i} / \sum_j e^{z_j}$. Penser à la loi de Boltzmann en physique statistique — l'exponentielle privilégie le plus grand $z_i$ et la division par la somme garantit que la masse totale fait 1.
Formule
$Z^{(1)} = W^{(1)} X + b^{(1)}$ ; $A^{(1)} = \sigma(Z^{(1)})$ avec $\sigma$ la sigmoïde ; $Z^{(2)} = W^{(2)} A^{(1)} + b^{(2)}$ ; $A^{(2)} = \text{softmax}(Z^{(2)}, \dim=0)$. Le résultat $A^{(2)} \in \mathbb{R}^{10 \times m}$ contient pour chaque exemple les 10 probabilités de classe. L'argument dim=0 indique que la normalisation softmax se fait sur la dimension des classes (les lignes), pas sur celle des exemples.
Piège
Le notebook utilise nn.Softmax(dim=0) Cell 9 parce que la convention features-en-lignes mise X en shape (784, N) et la sortie en shape (10, N). La somme doit être faite sur l'axe des classes (dim=0). Avec dim=1, on ferait la somme sur l'axe des exemples — chaque exemple aurait une 'fraction' du total, n'importe quoi. Le dim n'est pas optionnel.