Transformer Math · Authored by Sandeep Belgavi

Beam Search: Pruning an Explicit Probability Tree

Compare bounded beam search with exhaustive enumeration, inspect every retained and pruned candidate, and test EOS and length-normalization boundaries.

Experiment

Diagram labels stay full size. On narrow screens, scroll the diagram horizontally or focus it and use the arrow keys. Exact values follow below.

Paused · 0 / 60 teaching seconds

The timeline reveals computed states. Its sixty seconds describe presentation, not hardware speed.

Beam Search: Pruning an Explicit Probability Tree: experiment guide

This experiment searches a small conditional probability tree with the vocabulary A, B and EOS. It executes candidate expansion and pruning; the fixture is not a trained language model. Initially, A has probability 0.6 times one minus Initial EOS probability, B has probability 0.4 times that quantity, and EOS receives the selected probability directly. After A, the next-token probabilities for A, B and EOS are 0.45, 0.45 and 0.1. After B they are 0.05, 0.05 and 0.9. These rows sum to one. Only the last token selects the next row. The probability table does not learn, understand text or depend on a hidden random seed.

Worked details 1

Beam width sets the maximum number of retained candidates after each expansion, from one through eight. Maximum generated tokens bounds the search depth from one through six, including EOS when generated. Each active prefix expands to every positive-probability next token. A finished EOS sequence is carried forward unchanged and competes for the same beam slots. This explicit combined-pool policy differs from implementations that maintain separate alive and finished pools. Candidates are ranked by log joint probability divided by generated length raised to Length normalization alpha. Alpha zero ranks raw joint probabilities. Positive alpha can favor longer sequences because dividing a negative log probability by a larger number moves it toward zero. Length includes EOS. Equal scores use lexicographic token-path order, so A precedes B in exact ties.

Worked details 2

At the reset settings, beam width is two, maximum depth is three, alpha is zero and initial EOS probability is zero. At depth one the retained paths are A with probability 0.6 and B with probability 0.4. At depth two, B EOS has probability 0.36 and survives with A A at probability 0.27; the equally probable A B loses the tie. B EOS is carried through the next depth without another probability factor. It remains the winning completed sequence with probability 0.36. The table lists all candidates at the inspected depth, their parent, conditional probability, joint probability, log probability, normalized score and kept or pruned status. The tree draws retained nodes only; its blue active nodes and brown EOS nodes also have explicit token labels.

Worked details 3

Change Beam width to one and seek to the end. The search first discards B, then keeps A A and finally A A A, with probability 0.1215. The independently enumerated bounded reference still selects B EOS with probability 0.36. The displayed score gap is approximately 1.086190 natural-log units. This demonstrates a search failure caused by pruning, not a numerical error or an assertion that greedy search always fails. Set maximum depth to one: A is the highest-probability prefix but is labeled incomplete. The bounded reference compares both terminated sequences and prefixes truncated by the same depth limit, so it is not an oracle for unrestricted completed sequences.

Worked details 4

Reset, set Initial EOS probability to 0.5 and inspect the final depth. Alpha zero selects the one-token EOS path with probability 0.5. Change alpha to two: this fixture instead selects A B EOS with probability 0.1215 and normalized score approximately minus 0.234205. Its raw probability is lower, but the scoring objective changed. Initial EOS probability one removes both other root edges, ends after one expansion and produces probability one. Zero-probability edges are omitted instead of taking a logarithm of zero. Blank, out-of-range and fractional count inputs produce validation feedback and disable Play and Next depth; the model never silently rounds a requested beam width.

Worked details 5

Play reveals the computed depth history over sixty teaching seconds, even when the search terminates early. Pause freezes inspection, Next depth moves to the next retained level, and Seek pauses at a chosen time. Restart preserves parameters and returns to the root; Reset parameters restores the opening fixture. Speed changes presentation to quarter, normal or four times the teaching rate without changing any candidate score. Editing a model input rebuilds the trace and leaves it paused. The exhaustive comparison appears at the final inspected level. Its bounded enumeration is deliberately small enough to audit; it is not a practical strategy for a real language-model vocabulary. There is no KV cache, tokenizer, hardware timing, diversity penalty, minimum length or guarantee that a wider beam improves human-rated output quality.

Views are my own and do not represent any employer, past or present.