Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Hierarchical Softmax › The Problem: Expensive Normalization

Michael BrenndoerferUpdated February 8, 20265 min read

Language AI HandbookAll chapters

Part 4 of 67Word Embeddings

Chapter 22Hierarchical Softmax › The Problem: Expensive Normalization

Reading section 2 of 16 · The Problem: Expensive Normalization · 5 minSection 2 of 16 · 5 min

Section 2 of 16 · 5 min read

Hierarchical Softmax › The Problem: Expensive NormalizationOne section · Sections in this chapter2 / 16 ↓

Sections in this chapter16

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

To understand why we need hierarchical softmax, we must first appreciate how expensive standard softmax really is. For a center word $w_c$ with input embedding $\mathbf{h}$ , the probability of context word $w_j$ under standard softmax is:

P(w_j | w_c) = \frac{\exp(\mathbf{w}'_j \cdot \mathbf{h})}{\sum_{k=1}^{V} \exp(\mathbf{w}'_k \cdot \mathbf{h})}

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}$ : the input embedding of the center word $w_c$
$V$ : total vocabulary size

The denominator is called the partition function, or normalization constant. It ensures the output is a valid probability distribution, but computing it requires evaluating $\exp(\mathbf{w}'_k \cdot \mathbf{h})$ for every single word $k$ in the vocabulary, then summing all $V$ results. This is a full scan of the output embedding matrix, and it happens at every single training step.

To understand why this is so costly, consider the training scale that makes word embeddings useful. A typical word2vec training run processes two to ten billion word pairs. At each step, for each pair, we need to compute dot products with all $V$ output vectors to form the denominator. For a vocabulary of 100,000 words with embedding dimension 300, each softmax denominator requires 100,000 vector dot products of length 300, totaling 30 million multiplications and additions. Multiply by billions of training steps and the arithmetic quickly becomes untenable.

The problem is not just raw computation. The $O(V)$ cost also creates a bottleneck that prevents parallelism from helping as much as it otherwise would. Even with multiple GPUs, the bottleneck of the full vocabulary scan limits how fast training can proceed. Approximation methods that reduce this cost per step are therefore not mere optimizations: they are what makes large-scale word embedding training feasible at all.

In[4]:

Code

import time

import numpy as np


def standard_softmax_cost(vocab_size, embedding_dim, num_samples=1000):
    """Measure the average time for one standard softmax computation."""
    h = np.random.randn(embedding_dim)
    W_prime = np.random.randn(vocab_size, embedding_dim)

    start = time.time()
    for _ in range(num_samples):
        z = W_prime @ h
        exp_z = np.exp(z - np.max(z))
        probs = exp_z / np.sum(exp_z)
    elapsed = time.time() - start

    return elapsed / num_samples * 1000  # ms per step


vocab_sizes = [1_000, 10_000, 50_000, 100_000]
softmax_times = {V: standard_softmax_cost(V, 100) for V in vocab_sizes}

Out[5]:

Console

Standard Softmax Computation Time per Training Step:
------------------------------------------------
 Vocabulary Size      Time (ms)
------------------------------------------------
           1,000          0.024
          10,000          0.149
          50,000          0.620
         100,000          1.144

The timing confirms the linear scaling: doubling vocabulary size roughly doubles computation time. At 100,000 words, each training step is already expensive. Scale this up to billions of examples, and the arithmetic becomes sobering. We need an approach that reduces per-step cost without sacrificing the quality of learned embeddings.

Two main approaches emerged to solve this problem. Negative sampling, covered in a companion chapter, reframes training as binary classification and sidesteps normalization entirely. Hierarchical softmax takes a more principled route: it keeps the probabilistic interpretation but restructures the computation so that normalization becomes implicit rather than explicit, and local rather than global.

Out[6]:

Visualization

Log-log plot with two curves: red O(V) line and green O(log V) line, showing the widening gap. — Comparison of O(V) linear scaling (standard softmax, red) versus O(log V) logarithmic scaling (hierarchical softmax, green) across vocabulary sizes from 100 to 1,000,000. The logarithmic curve remains nearly flat while the linear curve grows without bound, illustrating the fundamental advantage of the tree-based decomposition.

Semi-log plot showing speedup factor growing from about 150x at 1000 words to over 6000x at 100,000 words. — Speedup factor of hierarchical softmax over standard softmax, defined as V divided by log2(V). At a vocabulary of 100,000 words, hierarchical softmax requires roughly 6,000 fewer operations per training step, translating directly into faster training on real corpora.

At 100,000 words, hierarchical softmax achieves thousands of times fewer operations per prediction. This is not a marginal improvement. It is the difference between a training run that completes in hours and one that takes months.