Hierarchical Softmax › A Worked Example: Tracing One Training Step
Language AI HandbookAll chapters Hierarchical Softmax › A Worked Example: Tracing One Training StepOne section · Sections in this chapter8 / 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
Let's trace through a complete training step to solidify the mechanics. We use a vocabulary of 12 words from the frequency distribution we built earlier, and predict "dog" as a context word given "the" as the center word.
The training step proceeds as follows. First, we look up the center word "the" in the input embedding matrix to get the context embedding . Then we look up "dog" in the Huffman path table to get its sequence of node indices and directions. We iterate through each node on the path, computing the dot product , the signed version , and the sigmoid. The loss is the sum of at each step. Gradients flow back: each node vector gets a small update, and the accumulated gradient is applied to the center word embedding.
import numpy as np
def sigmoid(x):
return np.where(x >= 0, 1 / (1 + np.exp(-x)), np.exp(x) / (1 + np.exp(x)))
class SimpleHierarchicalSoftmax:
"""Minimal hierarchical softmax implementation for demonstration."""
def __init__(self, word_freq: dict, embedding_dim: int):
self.dim = embedding_dim
# Build Huffman tree and extract paths
self._build_tree(word_freq)
# Initialize embeddings
np.random.seed(0)
vocab = list(word_freq.keys())
self.word_to_idx = {w: i for i, w in enumerate(vocab)}
n_words = len(vocab)
# Input embeddings (for center words)
self.W = np.random.randn(n_words, embedding_dim) * 0.01
# Node vectors (for internal nodes)
n_internal = n_words - 1
self.node_vecs = np.random.randn(n_internal, embedding_dim) * 0.01
def _build_tree(self, word_freq: dict):
"""Build Huffman tree and store paths as (node_indices, directions)."""
import heapq
# Build tree nodes
nodes = [
(freq, i, {"word": w, "left": None, "right": None, "idx": i})
for i, (w, freq) in enumerate(word_freq.items())
]
heapq.heapify(nodes)
next_idx = len(word_freq)
internal_nodes = {}
while len(nodes) > 1:
f1, _, n1 = heapq.heappop(nodes)
f2, _, n2 = heapq.heappop(nodes)
parent = {"word": None, "left": n1, "right": n2, "idx": next_idx}
internal_nodes[next_idx] = parent
heapq.heappush(nodes, (f1 + f2, next_idx, parent))
next_idx += 1
root = nodes[0][2]
# Extract paths for each word
self.paths = {}
def traverse(node, node_ids, directions):
if node["word"] is not None:
self.paths[node["word"]] = (node_ids, directions)
return
internal_idx = node["idx"] - len(word_freq)
traverse(node["left"], node_ids + [internal_idx], directions + [1])
traverse(
node["right"], node_ids + [internal_idx], directions + [-1]
)
traverse(root, [], [])
def compute_loss(self, center_word: str, context_word: str) -> float:
"""Compute HS loss for one (center, context) pair."""
h = self.W[self.word_to_idx[center_word]]
node_ids, directions = self.paths[context_word]
loss = 0.0
for nid, d in zip(node_ids, directions):
s = d * np.dot(self.node_vecs[nid], h)
loss += -np.log(float(sigmoid(s)) + 1e-12)
return loss
def train_step(self, center_word: str, context_word: str, lr: float = 0.1):
"""One SGD update for a (center, context) pair."""
center_idx = self.word_to_idx[center_word]
h = self.W[center_idx].copy()
node_ids, directions = self.paths[context_word]
grad_h = np.zeros(self.dim)
for nid, d in zip(node_ids, directions):
v = self.node_vecs[nid]
s = d * np.dot(v, h)
error = float(sigmoid(s)) - 1.0
grad_h += error * d * v
self.node_vecs[nid] -= lr * error * d * h
self.W[center_idx] -= lr * grad_h
return self.compute_loss(center_word, context_word)
# Create model and run training
hs_model = SimpleHierarchicalSoftmax(word_freq, embedding_dim=16)
center_word = "the"
context_word = "dog"
# Record loss before and after update
loss_before = hs_model.compute_loss(center_word, context_word)
hs_model.train_step(center_word, context_word, lr=0.1)
loss_after = hs_model.compute_loss(center_word, context_word)
# Run 100 steps to observe convergence
losses = [loss_before]
for _ in range(100):
l = hs_model.train_step(center_word, context_word, lr=0.05)
losses.append(l)Training step: ('the', 'dog')
Path length for 'dog': 5 nodes
Directions: ['right', 'left', 'left', 'right', 'left']
Loss before update: 3.4657
Loss after 1 step: 3.4654
Loss after 100 steps: 0.5013
Loss reduction: 85.5%
The training curve shows the expected pattern. The loss drops sharply in early steps as the model quickly learns the correct direction at each binary decision point. It then levels off as the binary classifiers become increasingly confident, with the error terms shrinking toward zero and reducing gradient magnitude accordingly. The final loss asymptotically approaches a lower bound determined by how well the context embedding can encode the information needed to follow the entire path.