Problem 105 · Tokenization & Sampling

Beam Search with Length Normalisation

HardAI engineering

01Problem

Given `step_fn(prefix) -> dict[token, logprob]`, implement `beam_search(step_fn, beam_width, max_len, eos, alpha)`. Keep the `beam_width` best live hypotheses per step; a hypothesis that emits `eos` moves to a finished list. Rank candidates by total log-probability (ties: lexicographic token order). At the end, return the best finished-or-live hypothesis by `logprob / len(seq) ** alpha` (alpha = 0 means raw log-probability).

02Key insight

Greedy decoding commits to the locally best token and can miss a higher-probability sequence whose first token was second-best; beam search keeps alternatives alive. But raw log-probability is a sum of negatives, so beams prefer short outputs; length normalisation (the alpha exponent) is the knob that corrects it, and the test shows the winner flipping. Finished hypotheses must leave the beam rather than occupy slots, or a single early EOS starves the search. Modern chat LLMs rarely use beam search because it produces bland, repetitive text, but it still matters for translation, ASR, and constrained structured generation.

03Solution

solution.pypython32 lines

from typing import Callable

Seq = tuple[str, ...]


def beam_search(step_fn: Callable[[Seq], dict[str, float]], beam_width: int, max_len: int,
                eos: str = "</s>", alpha: float = 0.0) -> tuple[Seq, float]:
    beams: list[tuple[Seq, float]] = [((), 0.0)]
    finished: list[tuple[Seq, float]] = []
    for _ in range(max_len):
        candidates = [
            (seq + (tok,), lp + tlp)
            for seq, lp in beams
            for tok, tlp in step_fn(seq).items()
        ]
        candidates.sort(key=lambda c: (-c[1], c[0]))
        beams = []
        for seq, lp in candidates:
            if seq[-1] == eos:
                finished.append((seq, lp))
            else:
                beams.append((seq, lp))
            if len(beams) == beam_width:
                break
        if not beams:
            break
    pool = finished + beams

    def score(c: tuple[Seq, float]) -> float:
        return c[1] / (len(c[0]) ** alpha) if alpha else c[1]

    return max(pool, key=score)

Try it first.

TimeO(max_len * B * V log(B * V))

SpaceO(B * max_len + finished)

tests.py · 4 passing

tests.pypython17 lines

import math
TABLE = {
    (): {"A": 0.4, "B": 0.3, "</s>": 0.3},
    ("A",): {"C": 0.5, "D": 0.5},
    ("B",): {"E": 0.9, "F": 0.1},
}
def step(prefix):
    dist = TABLE.get(prefix, {"</s>": 1.0})
    return {t: math.log(p) for t, p in dist.items()}
seq, lp = beam_search(step, 1, 5)
assert seq == ("A", "C", "</s>") and abs(math.exp(lp) - 0.2) < 1e-12
seq, lp = beam_search(step, 2, 5, alpha=0.0)
assert seq == ("</s>",) and abs(math.exp(lp) - 0.3) < 1e-12
seq, lp = beam_search(step, 2, 5, alpha=1.0)
assert seq == ("B", "E", "</s>") and abs(math.exp(lp) - 0.27) < 1e-12
seq, _ = beam_search(lambda p: {"x": math.log(0.9), "</s>": math.log(0.1)}, 2, 3)
assert seq in {("x", "x", "x"), ("</s>",)}

04Staff-level follow-up

On a real model each step is one batched forward pass over B beams sharing a prefix, so KV-cache reuse (copy-on-write or paged blocks forked per beam) dominates the cost. Add early stopping once the best finished score beats any live beam's optimistic bound, and diverse beam search or n-gram blocking when outputs repeat.