Hierarchical Softmax › Implementation: Complete Hierarchical Softmax Layer

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 HuffmanNode:
    def __init__(self, word=None, freq=0):
        self.word = word
        self.freq = freq
        self.left = None
        self.right = None
        self.index = None  # index into node_vecs array

    def __lt__(self, other):
        return self.freq < other.freq


def build_huffman_tree(word_freq: dict):
    """Build Huffman tree; return root and path lookup."""
    import heapq

    leaves = [HuffmanNode(word=w, freq=f) for w, f in word_freq.items()]
    heap = list(leaves)
    heapq.heapify(heap)

    internal_counter = [0]

    while len(heap) > 1:
        left = heapq.heappop(heap)
        right = heapq.heappop(heap)
        parent = HuffmanNode(freq=left.freq + right.freq)
        parent.index = internal_counter[0]
        internal_counter[0] += 1
        parent.left = left
        parent.right = right
        heapq.heappush(heap, parent)

    root = heap[0]

    # Assign indices to internal nodes via BFS
    from collections import deque

    queue = deque([root])
    idx = 0
    while queue:
        node = queue.popleft()
        if node.word is None:
            node.index = idx
            idx += 1
        if node.left:
            queue.append(node.left)
        if node.right:
            queue.append(node.right)

    # Extract paths
    paths = {}

    def extract_paths(node, node_indices, directions):
        if node.word is not None:
            paths[node.word] = (node_indices, directions)
            return
        extract_paths(node.left, node_indices + [node.index], directions + [1])
        extract_paths(
            node.right, node_indices + [node.index], directions + [-1]
        )

    extract_paths(root, [], [])
    n_internal = idx
    return root, paths, n_internal


class HierarchicalSoftmaxModel:
    """Skip-gram model with hierarchical softmax output layer."""

    def __init__(self, vocab: list, word_freq: dict, dim: int = 50):
        self.dim = dim
        self.vocab = vocab
        self.word_to_idx = {w: i for i, w in enumerate(vocab)}

        # Build Huffman tree
        _, self.paths, n_internal = build_huffman_tree(word_freq)

        # Learnable parameters
        np.random.seed(42)
        n_words = len(vocab)
        self.W_in = np.random.randn(n_words, dim) * 0.01  # input embeddings
        self.W_node = np.random.randn(n_internal, dim) * 0.01  # node vectors

    def forward_loss(self, center_word: str, context_word: str) -> float:
        """Compute loss for one training pair."""
        h = self.W_in[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 * float(np.dot(self.W_node[nid], h))
            loss -= np.log(float(sigmoid(s)) + 1e-10)
        return loss

    def train_step(
        self, center_word: str, context_word: str, lr: float = 0.025
    ):
        """SGD update for one (center, context) pair."""
        center_idx = self.word_to_idx[center_word]
        h = self.W_in[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.W_node[nid]
            s = d * np.dot(v, h)
            err = float(sigmoid(s)) - 1.0
            self.W_node[nid] -= lr * err * d * h
            grad_h += err * d * v

        self.W_in[center_idx] -= lr * grad_h

    def train(self, training_pairs: list, epochs: int = 5, lr: float = 0.025):
        """Train on a list of (center, context) pairs."""
        epoch_losses = []
        for epoch in range(epochs):
            np.random.shuffle(training_pairs)
            total_loss = sum(
                self.forward_loss(c, ctx) for c, ctx in training_pairs
            )
            epoch_losses.append(total_loss / len(training_pairs))
            for center, context in training_pairs:
                self.train_step(center, context, lr)
        return epoch_losses

# Build a small corpus and train
corpus = (
    "the cat sat on the mat the dog ran in the yard "
    "a cat and a dog are animals cats run dogs jump "
    "the quick brown fox jumps over the lazy dog "
    "happy cat happy dog sad cat sad run walk jump "
    "a evanescent sycophant ephemeral mellifluous cat dog run walk"
).split()

# Build vocabulary and training pairs (window size = 2)
word_counts = Counter(corpus)
vocab = list(word_counts.keys())
window_size = 2
training_pairs = []
for i, center in enumerate(corpus):
    for offset in range(-window_size, window_size + 1):
        if offset == 0:
            continue
        j = i + offset
        if 0 <= j < len(corpus):
            training_pairs.append((center, corpus[j]))

model = HierarchicalSoftmaxModel(vocab, word_counts, dim=16)
epoch_losses = model.train(training_pairs, epochs=15, lr=0.05)