Hierarchical Softmax: Efficient Word Probability Computation with Binary Trees

In the previous chapter, we encountered Skip-gram's computational bottleneck: the softmax normalization requires summing over all $V$ vocabulary words for every single training step. With vocabularies of 100,000+ words and billions of training examples, this $O(V)$ operation becomes prohibitively expensive. Training would take weeks or months.

Hierarchical softmax solves this problem. Instead of computing probabilities over a flat vocabulary, we organize words as leaves of a binary tree. Each word's probability is then computed as a product of binary decisions along the path from root to leaf. This reduces the complexity from $O(V)$ to $O(\log V)$, a dramatic improvement that makes training practical on large vocabularies.

This chapter explores hierarchical softmax from the ground up. We'll build the binary tree structure using Huffman coding, derive the path probability computation, work through the gradient updates, and implement a complete hierarchical softmax layer. By the end, you'll understand not just how this technique works, but why it's particularly well-suited for word embedding training.

Let's revisit why standard softmax is so expensive. For a center word $w_c$ with embedding $\mathbf{h}$, the probability of a context word $w_j$ is:

where:

• $P(w_j | w_c)$: probability of context word $w_j$ given center word $w_c$
• $\mathbf{w}'_j$: output embedding vector for word $w_j$
• $\mathbf{h}$: hidden layer representation (input embedding of center word)
• $V$: vocabulary size

The denominator, called the partition function or normalization constant, requires computing $V$ dot products and $V$ exponentials. Every training step needs this full computation, regardless of which context word we're predicting.

The scaling is $O(V)$: doubling vocabulary roughly doubles computation time. This linear scaling means a 100,000-word vocabulary requires 100 times more computation than a 1,000-word vocabulary. With billions of training examples, this becomes the dominant computational cost. We need a fundamentally different approach.

The visualization makes clear why hierarchical softmax matters: at realistic vocabulary sizes of 100,000+ words, we achieve thousands-fold speedups.

The key idea of hierarchical softmax is reorganizing the problem. Instead of asking "what's the probability of word $w_j$?" directly, we decompose it into a sequence of binary questions: "at this node, should we go left or right?"

Consider a simple example with 8 words. A balanced binary tree has depth 3 ($\log_2 8 = 3$). To compute the probability of any word, we make exactly 3 binary decisions, each requiring just one sigmoid computation. Compare this to standard softmax, which computes 8 exponentials and sums them.

The visualization shows the core idea: computing $P(\text{dog}|\mathbf{h})$ requires only 3 binary decisions instead of 8 softmax terms. For a vocabulary of 100,000 words, a balanced tree has depth $\log_2(100,000) \approx 17$, so we need only 17 operations instead of 100,000.

With the tree structure in place, we can now develop the mathematical machinery for computing word probabilities. This section builds the formula step by step, starting from the intuition of tree traversal and arriving at a compact expression that captures hierarchical softmax in its entirety.

Imagine standing at the root of the tree with a specific word in mind, say "dog." To reach this word, you must navigate through a series of decision points: at each internal node, you choose either the left branch or the right branch. The path you take is unique to each word, determined by its position in the tree.

We formalize this journey as a sequence of nodes: $n(w, 1), n(w, 2), \ldots, n(w, L(w))$, where:

• $n(w, 1)$ is always the root node
• $n(w, L(w))$ is the leaf node representing word $w$
• $L(w)$ is the total number of nodes along the path (path length)

The first $L(w) - 1$ nodes are internal nodes where decisions are made. The final node is the word itself, the destination of our journey.

Here's the key probabilistic insight: the probability of reaching a word equals the probability of making all the correct turns along its path. If reaching "dog" requires going left, then right, then right, we need to compute the probability of each of these three decisions and multiply them together.

But how does the model "decide" which direction to go at each node? This is where the learned vectors come in.

Each internal node $n$ has an associated vector $\mathbf{v}_n$ of the same dimension as the word embeddings. When we arrive at a node, we compare this node vector with the context embedding $\mathbf{h}$ (the embedding of the center word in Skip-gram). The dot product $\mathbf{v}_n \cdot \mathbf{h}$ measures how "compatible" this node is with the current context:

• A large positive dot product suggests the context "wants" to go left
• A large negative dot product suggests the context "wants" to go right
• A dot product near zero indicates uncertainty

To convert this compatibility score into a probability, we apply the sigmoid function:

where:

• $\sigma(\cdot)$: the sigmoid function, which maps any real number to the interval $(0, 1)$
• $\mathbf{v}_n$: the learned vector associated with internal node $n$
• $\mathbf{h}$: the context embedding (input embedding of the center word)
• $\mathbf{v}_n \cdot \mathbf{h}$: the dot product measuring compatibility between the node and context

This sigmoid output gives us the probability of going left at node $n$. The probability of going right is its complement: $1 - \sigma(\mathbf{v}_n \cdot \mathbf{h})$. Notice that these two probabilities always sum to 1, which is what guarantees the tree produces a valid probability distribution over all words.

Now we can assemble these pieces into the complete formula. The probability of word $w$ given context embedding $\mathbf{h}$ is the product of all binary decisions along the path:

where:

• $P(w | \mathbf{h})$: probability of word $w$ given the context embedding $\mathbf{h}$
• $L(w)$: the total number of nodes along the path from root to word $w$ (including both the root and the leaf)
• $n(w, j)$: the $j$-th node along the path to word $w$
• $\text{left}(n(w,j))$: the left child of node $n(w,j)$
• $[\![ \cdot ]\!]$: the Iverson bracket, which equals $+1$ if the condition inside is true and $-1$ if false
• $\mathbf{v}_{n(w,j)}$: the learned vector at the $j$-th internal node on the path
• $\mathbf{h}$: the context embedding (center word's input embedding)

This formula looks intimidating, but it encapsulates a simple idea: multiply together the probability of each correct turn. The Iverson bracket $[\![ n(w, j+1) = \text{left}(n(w,j)) ]\!]$ equals $+1$ if we went left at node $j$ (meaning the next node is the left child) and $-1$ if we went right.

Why does this encoding work? Consider the two cases:

• Going left ($[\![\cdot]\!] = +1$): We compute $\sigma(\mathbf{v}_n \cdot \mathbf{h})$, which is high when the dot product is positive
• Going right ($[\![\cdot]\!] = -1$): We compute $\sigma(-\mathbf{v}_n \cdot \mathbf{h}) = 1 - \sigma(\mathbf{v}_n \cdot \mathbf{h})$, which is high when the dot product is negative

This encoding allows us to express both directions with a single, unified formula.

The formula above is mathematically precise but notationally heavy for practical use. Let's introduce cleaner notation that we'll use for the rest of this chapter.

For each step $j$ along the path to word $w$, define:

• $d_j \in \{-1, +1\}$: the direction encoding at node $j$ ($+1$ = left, $-1$ = right)
• $\mathbf{v}_j$: the learned vector at the $j$-th internal node on the path

With this notation, the path probability simplifies to:

This is the central formula of hierarchical softmax. Let's unpack each component:

The key property of this formula is that it reduces computing a probability over $V$ words to computing just $L(w) - 1$ sigmoid operations, typically around $\log_2 V$ for a balanced tree.

The output reveals several important insights about how hierarchical softmax works in practice:

1. Probability accumulation: Each step multiplies into the running probability. Starting at 1.0, we progressively narrow down to the final word probability.

2. Direction matters: The d * dot term flips the sign appropriately. When going right (d = -1), a negative dot product becomes positive after multiplication, yielding a high sigmoid value.

3. Confidence varies by node: Step probabilities near 0.5 indicate the model is uncertain about the direction; values near 0 or 1 indicate high confidence.

4. Final probability can be small: With three multiplications, even reasonable step probabilities (around 0.7) compound to give a final probability around 0.35. For longer paths in large vocabularies, probabilities become very small, but that's expected when distributing probability mass over 100,000+ words.

A critical property that makes hierarchical softmax mathematically valid: the probabilities over all words sum to exactly 1. This isn't immediately obvious from the formula, but it follows directly from the tree structure.

The key insight is that the tree creates a complete partition of the probability space. At each internal node, the probability of going left plus the probability of going right equals 1. Since every word is reachable by exactly one path, and the paths collectively cover all possible ways to traverse the tree, the probabilities must sum to 1.

Let's verify this empirically:

The sum equals 1.0 (within floating-point precision), confirming our theoretical claim that the tree structure produces a valid probability distribution. This result has important practical implications:

• No normalization needed: Unlike standard softmax, which requires summing over all $V$ words to normalize, hierarchical softmax produces valid probabilities automatically.
• Computational savings: We compute only $O(\log V)$ operations per word, not $O(V)$.
• Proper probability model: The output is a true probability distribution, not an approximation. This matters for applications that need calibrated probabilities.

The visual bar chart also reveals something interesting: the word probabilities vary significantly even with random vectors. Once trained, these probabilities will reflect the actual word co-occurrence patterns in the training data.

So far we've assumed a balanced binary tree where every word has the same path length: $\log_2 V$ nodes. But this ignores a fundamental property of natural language: word frequencies are extremely skewed. In English, "the" appears about 7% of the time, while "serendipity" might appear once per million words.

This skewness creates an opportunity for optimization. If we could assign shorter paths to frequent words and longer paths to rare words, we'd reduce the average number of computations per training example. Since training involves billions of word predictions, even small savings per word compound into massive speedups.

The left plot shows the characteristic power-law relationship of Zipf's law on a log-log scale. The right plot reveals just how concentrated word usage is: roughly 50% of all text comes from just the top ~100 words, and 90% comes from less than ~500 words. This means the vast majority of training examples involve a small set of frequent words, exactly the words Huffman coding will place at shallow tree depths.

Enter Huffman coding, a classic algorithm from information theory that solves exactly this problem.

The Huffman algorithm is surprisingly simple given its optimality guarantees:

1. Initialize: Create a leaf node for each word, weighted by its frequency
2. Merge: Repeatedly combine the two lowest-weight nodes into a new internal node
3. Propagate: The new node's weight equals the sum of its children's weights
4. Repeat: Continue until only one node remains, which becomes the root

The output reveals the Huffman tree's structure:

• Frequency-depth relationship: "the" (most frequent, frequency 1000) has the shortest code, while "from" (least frequent, frequency 62) has the longest. This inverse relationship is the hallmark of Huffman coding.

• Binary codes as paths: Each code represents the path from root to leaf. A code of "00" means "go left twice," while "110" means "go left, then right twice."

• Computational savings: The weighted average path length is less than the balanced tree depth. This gap represents the efficiency gain: we perform fewer operations on common words that dominate training.

To understand the impact of Huffman coding, consider the expected number of operations per training word:

where:

• $\mathbb{E}[\text{path length}]$: expected number of binary decisions per word
• $P(w)$: probability of word $w$ appearing as a training target
• $L(w)$: path length (number of nodes) for word $w$ in the tree

This formula reveals why Huffman coding is so effective for natural language:

• Frequent words dominate: In a typical corpus, "the," "of," "and," "to," and "a" might account for 10-15% of all words. If these words have short paths (2-3 nodes), they contribute very little to the expected path length.

• Rare words are rarely seen: Words like "serendipity" or "ephemeral" might appear once per million words. Even if they have long paths (10+ nodes), their contribution to the expected path length is negligible.

• Zipf's law amplifies savings: Natural language follows Zipf's law, where word frequency is inversely proportional to rank. This extremely skewed distribution means Huffman coding provides substantial savings compared to a balanced tree.

The inverse relationship between word frequency and Huffman path length is clear from the data:

With the path probability formula and Huffman tree structure in place, we can now formalize the training objective. Recall that in standard Skip-gram, we maximize the log-probability $\log P(w_o | w_c)$ where $w_o$ is a context word and $w_c$ is the center word.

With hierarchical softmax, this log-probability decomposes cleanly. Since the word probability is a product of sigmoid terms, the log-probability becomes a sum:

where:

• $\log P(w_o | w_c)$: the log-probability of context word $w_o$ given center word $w_c$ (this is what we maximize during training)
• $\mathbf{h} = \mathbf{W}[w_c]$: the center word's embedding from the input embedding matrix $\mathbf{W}$
• $L(w_o)$: the path length to word $w_o$ (number of nodes from root to leaf, inclusive)
• $\mathbf{v}_j$: the learned vector at the $j$-th internal node on the path
• $d_j \in \{-1, +1\}$: the direction encoding at node $j$ ($+1$ for left, $-1$ for right)
• $\sigma(\cdot)$: the sigmoid function
• $\log \sigma(d_j \cdot \mathbf{v}_j \cdot \mathbf{h})$: the log-probability of taking the correct direction at node $j$

The key insight is that the log of a product becomes a sum: $\log \prod_j \sigma(\cdot) = \sum_j \log \sigma(\cdot)$. Each term in this sum represents a single binary decision. Maximizing the objective means teaching the model to make correct navigation decisions at every node along every path. When the model correctly predicts "go left" (high $\sigma$ when $d_j = +1$) or "go right" (high $\sigma$ when $d_j = -1$), the log-probability is close to zero. Incorrect predictions yield large negative values.

Let's compute this loss for a concrete example:

The output illustrates several key points about the loss function:

• Per-node contributions: Each internal node adds to the log-probability. Values close to 0 (like -0.1) indicate high confidence in the correct direction; more negative values (like -0.8) indicate uncertainty or incorrect predictions.

• Additive structure: Unlike the multiplicative path probability, the log-probability is additive. This is numerically more stable and leads to cleaner gradient formulas.

• Loss interpretation: The final loss is the negative log-probability. Lower loss means higher probability of the correct context word, exactly what we want during training.

• Scale of loss values: With random initialization, losses are typically in the range 2-5. After training, losses drop significantly as the model learns to navigate the tree correctly.

With the forward pass formula established, we now turn to training: how do we update the parameters to make the model assign higher probability to correct context words? This requires computing gradients of the loss with respect to two sets of parameters:

1. Node vectors $\mathbf{v}_j$: Each internal node has a learned vector that determines binary decisions
2. Input embedding $\mathbf{h}$: The center word's embedding, which gets passed back to update the embedding matrix

The gradient derivation reveals a structure that makes hierarchical softmax efficient to train.

We'll derive the gradients step by step, starting from a single node's contribution and then assembling the full path gradient.

Step 1: The single-node objective

At each node $j$ on the path, the model makes a binary decision. The contribution to the log-likelihood from this node is:

where:

• $L_j$: the log-probability contribution from node $j$ (we want to maximize this)
• $d_j \in \{-1, +1\}$: the direction encoding ($+1$ for left, $-1$ for right)
• $\mathbf{v}_j$: the learned vector at node $j$
• $\mathbf{h}$: the context embedding (center word's input embedding)
• $\sigma(\cdot)$: the sigmoid function

Our goal is to find $\frac{\partial L_j}{\partial \mathbf{v}_j}$ (how to update the node vector) and $\frac{\partial L_j}{\partial \mathbf{h}}$ (how to update the center word embedding).

Step 2: The sigmoid derivative identity

Let $z = d_j \cdot \mathbf{v}_j \cdot \mathbf{h}$. To find the derivative, we apply the chain rule to $\log \sigma(z)$:

where:

• $\sigma'(z)$: the derivative of the sigmoid function, which equals $\sigma(z)(1 - \sigma(z))$
• $\sigma(z)$: the sigmoid function evaluated at $z$

The first equality uses the chain rule for logarithms: $\frac{d}{dz}\log f(z) = \frac{f'(z)}{f(z)}$. The second equality substitutes the well-known sigmoid derivative identity $\sigma'(z) = \sigma(z)(1 - \sigma(z))$. The final simplification cancels $\sigma(z)$ from numerator and denominator.

Step 3: Applying the chain rule

Since $z = d_j \cdot \mathbf{v}_j \cdot \mathbf{h}$ depends on both $\mathbf{v}_j$ and $\mathbf{h}$, we apply the chain rule. First, we write down what we computed in Step 2:

Next, we need the partial derivatives of $z$ with respect to the vectors. Since $z = d_j \cdot (\mathbf{v}_j \cdot \mathbf{h})$ is a scalar formed by multiplying the direction encoding $d_j$ with the dot product of two vectors:

• $\frac{\partial z}{\partial \mathbf{v}_j} = d_j \cdot \mathbf{h}$ (the gradient of a dot product $\mathbf{v}_j \cdot \mathbf{h}$ with respect to $\mathbf{v}_j$ is $\mathbf{h}$)
• $\frac{\partial z}{\partial \mathbf{h}} = d_j \cdot \mathbf{v}_j$ (the gradient of a dot product $\mathbf{v}_j \cdot \mathbf{h}$ with respect to $\mathbf{h}$ is $\mathbf{v}_j$)

Combining via the chain rule $\frac{\partial L_j}{\partial \mathbf{v}_j} = \frac{\partial L_j}{\partial z} \cdot \frac{\partial z}{\partial \mathbf{v}_j}$:

where:

• $\frac{\partial L_j}{\partial \mathbf{v}_j}$: gradient of the log-likelihood with respect to the node vector (a vector of the same dimension as $\mathbf{v}_j$)
• $\frac{\partial L_j}{\partial \mathbf{h}}$: gradient with respect to the context embedding

Step 4: Converting to loss gradients

Since we minimize the negative log-likelihood (loss $= -L$), we negate these gradients. Negating flips the sign of $(1 - \sigma(\cdot))$ to $(\sigma(\cdot) - 1)$:

where:

• $\nabla_{\mathbf{v}_j} \text{Loss}$: gradient of the loss with respect to node vector $\mathbf{v}_j$
• $\nabla_{\mathbf{h}} \text{Loss}_j$: contribution to the embedding gradient from node $j$
• $\sigma(d_j \cdot \mathbf{v}_j \cdot \mathbf{h})$: the probability the model assigns to taking direction $d_j$ at node $j$

Step 5: The error signal interpretation

The term $\sigma(d_j \cdot \mathbf{v}_j \cdot \mathbf{h}) - 1$ has a natural interpretation as an error signal:

where:

• $e_j$: the error signal at node $j$, measuring how incorrect the model's prediction is
• $\sigma(d_j \cdot \mathbf{v}_j \cdot \mathbf{h})$: the probability the model assigns to taking direction $d_j$

This error signal $e_j$ lies in the range $[-1, 0]$:

With this notation, the gradients simplify to:

where:

• $\nabla_{\mathbf{v}_j} \text{Loss}$: gradient for updating node vector $\mathbf{v}_j$
• $\nabla_{\mathbf{h}} \text{Loss}$: total gradient for updating the center word embedding $\mathbf{h}$
• $L$: the path length (number of nodes from root to leaf)
• $d_j$: direction encoding at node $j$ ($+1$ for left, $-1$ for right)
• $e_j$: error signal at node $j$ (ranges from $-1$ to $0$)

Notice the structure: the embedding gradient is the sum of contributions from all nodes along the path. Each contribution is the node vector scaled by its error signal and direction encoding.

Error signals near 0 indicate the model is confident in the correct direction (small updates needed), while values near -1 indicate confidence in the wrong direction (large corrective updates). The gradient computation reveals several key insights about how the model learns:

• Node-specific updates: Each node vector $\mathbf{v}_j$ receives a gradient proportional to the error at that specific node. Nodes that already make correct predictions receive small updates.

• Accumulated embedding gradient: The center word embedding $\mathbf{h}$ receives contributions from all nodes along the path. This is analogous to backpropagation through time in RNNs: the embedding must learn to navigate the entire path correctly.

• Error-weighted learning: The magnitude of updates is automatically scaled by confidence. Uncertain predictions ($e_j \approx -0.5$) receive moderate updates, while confident wrong predictions ($e_j \approx -1$) receive the largest updates.

The visualization shows a key signature of successful training: error signals shift from being uniformly distributed (random guessing) to being concentrated near zero (confident correct predictions). This shift corresponds to decreasing loss values during training.

The following diagram illustrates how gradients flow through the tree structure:

With the mathematical foundations established, we can now implement a complete hierarchical softmax layer. This implementation brings together all the concepts we've developed:

• Huffman tree construction from word frequencies
• Path lookup for efficient word-to-path mapping
• Forward pass computing $P(w|\mathbf{h}) = \prod_j \sigma(d_j \cdot \mathbf{v}_j \cdot \mathbf{h})$
• Backward pass computing $\nabla_{\mathbf{v}_j}$ and $\nabla_{\mathbf{h}}$ using the error signals
• Parameter updates using stochastic gradient descent

The code below is a complete, working implementation that you can use to train word embeddings:

With the HierarchicalSoftmax class implemented, let's train word embeddings on a synthetic dataset. This demonstrates the complete training loop:

1. Forward pass: Compute the loss for each (center, context) pair
2. Backward pass: Compute gradients for embeddings and node vectors
3. Update: Apply gradient descent to all parameters

We'll use Zipf-distributed word frequencies to simulate realistic vocabulary statistics: