Hierarchical Softmax › Huffman Coding: Optimizing the Tree
Language AI HandbookAll chapters Hierarchical Softmax › Huffman Coding: Optimizing the TreeOne section · Sections in this chapter4 / 16 ↓
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
Not all trees are created equal. A randomly organized binary tree or a perfectly balanced one assigns paths of similar length to all words, giving every word the same computational cost per prediction. But natural language has a vastly skewed word frequency distribution: the word "the" appears billions of times in any large corpus, while "evanescent" might appear thousands of times and "mellifluous" perhaps hundreds. If we assign "the" a short path and "evanescent" a long path, frequent words get cheaper predictions on average, and the total computational cost weighted by training frequency drops substantially.
Huffman coding constructs precisely this kind of frequency-optimized tree. The algorithm builds the tree from the bottom up, starting with all individual words and repeatedly merging the two lowest-frequency nodes into a parent. Because low-frequency nodes are always merged first, high-frequency words end up closest to the root and receive the shortest paths.
Huffman coding is a greedy algorithm for constructing an optimal prefix-free binary tree. Given symbol frequencies, it repeatedly merges the two lowest-frequency nodes into a combined parent node, assigning that parent the sum of their frequencies. The resulting tree minimizes the expected code length when each symbol's code is its root-to-leaf path.
The algorithm is straightforward to implement:
- Place each word (with its frequency) into a priority queue ordered by frequency ascending.
- Extract the two lowest-frequency entries.
- Create a new parent node whose frequency is the sum of the two.
- Insert the parent node back into the queue.
- Repeat until only one node remains. That node is the root.
The reason this greedy approach is optimal deserves a moment of reflection. When we merge the two least-frequent nodes at each step, we are pushing them one level deeper in the tree. This is always safe because they are the least-frequent words, so they will be visited least often during training. Any other pair of nodes we could have chosen to merge would have resulted in a worse expected path length, because we would be pushing more frequently needed words deeper. The greedy strategy of always merging the two cheapest nodes turns out to produce the globally optimal tree.
The final expected path length under Huffman coding equals the entropy of the word frequency distribution, which is the theoretical minimum for any prefix-free binary tree. No other arrangement can achieve a shorter average path length. This is a well-known result from information theory: Huffman codes achieve Shannon's lower bound on average code length.
import heapq
from dataclasses import dataclass
from typing import Optional
@dataclass
class HuffmanNode:
"""A node in the Huffman tree."""
freq: int
word: Optional[str] = None
left: Optional["HuffmanNode"] = None
right: Optional["HuffmanNode"] = None
# Comparison uses freq only
def __lt__(self, other):
return self.freq < other.freq
def build_huffman_tree(word_freq: dict) -> HuffmanNode:
"""Build a Huffman tree from a word frequency dictionary."""
heap = [HuffmanNode(freq=f, word=w) for w, f in word_freq.items()]
heapq.heapify(heap)
while len(heap) > 1:
left = heapq.heappop(heap)
right = heapq.heappop(heap)
parent = HuffmanNode(
freq=left.freq + right.freq,
left=left,
right=right,
)
heapq.heappush(heap, parent)
return heap[0] # root
def get_paths(root: HuffmanNode) -> dict:
"""Extract root-to-leaf paths for all words."""
paths = {}
def traverse(node, path, directions):
if node.word is not None: # leaf
paths[node.word] = {
"path_length": len(directions),
"directions": directions,
}
return
if node.left:
traverse(node.left, path + [node], directions + [1])
if node.right:
traverse(node.right, path + [node], directions + [-1])
traverse(root, [], [])
return paths
# Build a small example vocabulary with realistic frequency skew
word_freq = {
"the": 5000,
"a": 4000,
"is": 3000,
"in": 2500,
"cat": 500,
"dog": 450,
"run": 300,
"jump": 200,
"evanescent": 20,
"sycophant": 15,
"ephemeral": 10,
"mellifluous": 5,
}
root = build_huffman_tree(word_freq)
paths = get_paths(root)Huffman Tree Path Lengths:
---------------------------------------------
Word Frequency Path Length
---------------------------------------------
is 3,000 2
a 4,000 2
the 5,000 2
in 2,500 3
run 300 5
dog 450 5
cat 500 5
jump 200 6
evanescent 20 7
sycophant 15 8
mellifluous 5 9
ephemeral 10 9
Expected path length: 2.46 nodes
Entropy lower bound: 2.42 bits
Balanced tree depth: 4.00 levelsHuffman coding assigns short paths to frequent words and long paths to rare words. The expected path length is close to the entropy of the frequency distribution, confirming the tree is nearly optimal. Frequent words like "the" and "a" get paths of length 2 to 3, while rare words like "mellifluous" get paths of 4 to 5. The average cost per training step is therefore much lower than a balanced tree would provide.
The connection to entropy is worth pausing on. Shannon entropy measures the average uncertainty in a random variable. A word frequency distribution with high entropy (many words with similar frequencies) will require a deeper Huffman tree because no words can be given especially short paths without penalizing others. A low-entropy distribution (dominated by a few very frequent words) allows the Huffman tree to put those words near the root and push the rare words far down, achieving dramatic savings in expected path length. Natural language frequency distributions are extremely low-entropy due to Zipf's law, which makes Huffman coding especially effective.
