Priority queues and Dijkstra to beam search

The decoder's beam is a priority queue of half-finished sentences - the exact frontier Dijkstra expands, only scored by log-probability instead of distance.

In Shortest paths: Dijkstra you watched Dijkstra grow a shortest-path tree by repeatedly pulling the closest unsettled node off a priority queue and expanding it - best-first search. Beam search, the workhorse decoder behind translation and captioning, is the very same machine pointed at a different graph: its nodes are half-written sentences, its edge weights are per-token log-probabilities, and its heap keeps the $k$ highest-scoring partial sequences alive at every step. Learn one and you have learned the other; the only real differences are what a node is and the fact that the beam is bounded to width $k$ instead of holding the whole frontier.

The engine you already own: a heap-backed frontier

Every one of these algorithms is a loop over the same two moves: pop the most promising item from a priority queue, then push its expansions back in. Dijkstra pops the min-cost node and pushes relaxed neighbours; beam search pops (surfaces) the top-$k$ sequences and pushes each one's next-token children. The heap is what makes “which candidate is best?” cheap enough to ask thousands of times per sentence. Here is the translation, term for term - nothing on the right is a metaphor.

The one honest difference: Dijkstra keeps every reachable node on the heap and is therefore exact. Beam search deliberately throws away all but the top $k$ at each step to stay affordable - the sequence graph has $V^L$ paths ($V$ = vocabulary, $L$ = length), far too many to settle exactly. That truncation is what makes beam search a fast heuristic rather than a guaranteed optimum.

Scoring a path: add up the log-probabilities

A model gives you $p_{\theta }(y_i\mid y_{<i})$ - the probability of each next token given the prefix. The probability of a whole sequence is the product of these, but a decoder scores it as the sum of log-probabilities:

Three reasons for the log. Multiplying hundreds of probabilities underflows to zero in floating point; $\log$ is monotonic, so it never changes the ranking; and it turns a product into a sum, converting “probability of a path” into an additive cost - exactly the additive edge weights Dijkstra accumulates along a route. Negate the log-probs and “more probable” becomes “less costly,” and beam search is Dijkstra with a fixed-width frontier.

One subtlety follows immediately. Because every $\log p\le 0$, longer paths keep adding negative terms, so a raw score always favours shorter sequences. Beam search corrects this with length normalization, dividing by the length raised to a power $\alpha \in [0.6,0.7]$:

Within a single step every candidate on the beam has the same length, so this rescaling changes nothing there - it only matters when you finally compare finished sentences of different lengths (once an end-of-sequence token can fire).

Beam width $k$: from greedy to (almost) exhaustive

The demo below decodes a toy language model, three candidate tokens per step, three steps deep. Drag the beam width $k$ and step through the levels. Each step pushes every surviving beam's children into the priority queue on the right; the heap sorts them by score and surfaces the top $k$; everything below the cut is pruned and its branch fades in the tree. Watch how $k=1$ is plain greedy decoding and larger $k$ widens toward exhaustive search.

import heapq, math

def beam_search(next_dist, k, steps):
    # each beam is (cumulative log-prob score, list of tokens)
    beams = [(0.0, ["<s>"])]
    for _ in range(steps):
        frontier = []                              # the priority queue for this level
        for score, seq in beams:
            for token, p in next_dist(seq):        # PUSH every child
                heapq.heappush(frontier, (score + math.log(p), seq + [token]))
        beams = heapq.nlargest(k, frontier)        # SURFACE top-k, prune the rest
    return max(beams)                              # best-scoring complete sequence

# Same loop as Dijkstra: swap `next_dist` for graph edges and `log(p)`
# for -weight, drop the top-k cap, and you are back to shortest paths.

Check yourself