Hierarchical Softmax › Hierarchical Softmax

The skip-gram model learns powerful word representations, but it carries a serious computational burden: the softmax denominator. Every training step requires summing over every word in the vocabulary to normalize probabilities. For a vocabulary of 100,000 words, that means 100,000 dot products and exponentials for a single prediction. With billions of training examples, standard softmax makes large-scale word embedding training impractically slow.

Hierarchical softmax solves this problem by reorganizing how we think about word prediction. Instead of choosing among all words at once, we arrange words as the leaves of a binary tree and decompose each prediction into a sequence of binary decisions. To predict any word, we trace a path from the root to that word's leaf, multiplying the probabilities of turning left or right at each internal node. This reduces per-step complexity from $O(V)$ to $O(\log V)$, a transformation that makes training on large vocabularies tractable.

The elegance of hierarchical softmax lies in what it preserves alongside what it eliminates. Standard softmax computes a true probability distribution: all word probabilities sum to exactly 1. Hierarchical softmax achieves the same guarantee through a completely different mechanism. By ensuring that the left and right probabilities at every internal node sum to 1, the tree guarantees by induction that the total probability mass across all leaves also sums to 1. You get the mathematical correctness of a full probability distribution without paying the $O(V)$ cost to compute it.

This chapter develops hierarchical softmax from the ground up. We begin with the computational problem that motivates it, then build intuition for how a binary tree decomposition works, derive the path probability formula and its gradients, implement a complete hierarchical softmax layer in NumPy, and finally compare the approach against negative sampling to understand when each technique is the right choice.