Beam Search
Beam search is a heuristic search algorithm that explores a graph by expanding the most promising nodes in a limited set. It is an optimization of best-first search that reduces its memory consumption by maintaining a fixed-width frontier (the beam width, B) at every depth level T, pruning sub-optimal paths to trade exhaustive optimality for deterministic computational efficiency.
HOW-TO IMPLEMENT & CONFIGURE
- Python 3.10+
- PyTorch 2.0+ or NumPy 1.24+
- Typing and dataclasses standard libraries
Step 1
Define the foundational data structures required to track active hypotheses, sequence tokens, and cumulative log-probabilities across decoding steps.
from dataclasses import dataclass
from typing import List
import torch
@dataclass
class Hypothesis:
tokens: List[int]
score: float
def __lt__(self, other: 'Hypothesis') -> bool:
return self.score > other.score # Max-heap behavior for sorting highest scoresStep 2
Implement the iterative beam expansion, score accumulation, and top-B pruning loop, incorporating sequence termination handling.
def beam_search_decode(model, start_token: int, end_token: int, max_length: int, beam_width: int, vocab_size: int) -> List[int]:
# Initialize beam with a single hypothesis containing the start token
current_beams = [Hypothesis(tokens=[start_token], score=0.0)]
completed_beams = []
for _ in range(max_length):
next_beams = []
for hyp in current_beams:
if hyp.tokens[-1] == end_token:
completed_beams.append(hyp)
continue
# Simulate model forward pass yielding log-probabilities for vocabulary
tensor_input = torch.tensor([hyp.tokens], dtype=torch.long)
logits = model(tensor_input) # Shape: (1, vocab_size)
log_probs = torch.log_softmax(logits, dim=-1).squeeze(0)
topk_probs, topk_indices = torch.topk(log_probs, beam_width)
for i in range(beam_width):
token = topk_indices[i].item()
score = hyp.score + topk_probs[i].item()
next_beams.append(Hypothesis(tokens=hyp.tokens + [token], score=score))
if not current_beams:
break
# Sort and select top B beams
next_beams.sort(key=lambda x: x.score, reverse=True)
current_beams = next_beams[:beam_width]
if completed_beams:
completed_beams.sort(key=lambda x: x.score, reverse=True)
return completed_beams[0].tokens
return current_beams[0].tokensStep 3
Construct a deterministic mock model to validate that the beam search correctly selects optimal paths and respects length and termination bounds.
import unittest
class MockModel:
def __call__(self, x):
# Mock static logits favoring deterministic sequence progression
vocab_size = 10
logits = torch.zeros((1, vocab_size))
next_token = (x.shape[1]) % vocab_size
logits[0, next_token] = 10.0
return logits
class TestBeamSearch(unittest.TestCase):
def test_beam_search_execution(self):
model = MockModel()
result = beam_search_decode(model, start_token=1, end_token=5, max_length=5, beam_width=2, vocab_size=10)
self.assertIsInstance(result, list)
self.assertGreater(len(result), 0)
if __name__ == '__main__':
unittest.main()Step 4
Apply length penalty scaling to the scoring function to prevent systemic degradation toward short sequence outputs.
def apply_length_penalty(score: float, length: int, alpha: float = 0.7) -> float:
# Standard GNMT (Google Neural Machine Translation) length penalty
penalty = ((5.0 + length) ** alpha) / ((5.0 + 1.0) ** alpha)
return score / penaltyMathematical Rigor & Formal Definition
Let G = (V, E) be a directed acyclic graph representing a sequence generation space, where V represents partial hypotheses and E represents valid state transitions driven by a probabilistic scoring function P(y_t | y_<t). Exhaustive breadth-first search grows the frontier exponentially at O(V^T), where V is the vocabulary size and T is the target sequence length. Beam search constrains the frontier size at each step t to a constant upper bound B, known as the beam width. At each decoding step, the algorithm generates all valid single-step extensions for each of the B active hypotheses, evaluates their cumulative score, and applies a top-k selection operator to retain only the top B highest-scoring hypotheses, discarding the rest. This pruning strategy transforms the exponential time complexity into O(T * B * |V|), and spatial complexity from O(V^T) to O(B * T).
Algorithmic Trade-offs & Failure Modes
While beam search drastically reduces computational overhead, it introduces notable limitations. First, it is fundamentally greedy at the macro-level; discarding a path with a low intermediate score can lead to missed optimal global solutions if a high-reward state is preceded by a low-probability transition. Second, length bias frequently impacts decoding, as cumulative log-probabilities naturally favor shorter sequences. Modern implementations mitigate this via length normalization heuristics or coverage penalties. Third, beam search is susceptible to repetition loops, where the model generates cyclical token sequences that continuously score within the top-B threshold. This requires strict repetition penalties or unigram/ngram blocking mechanisms during sequence generation.
📚 STANDARDS & ACADEMIC RFC CITATIONS
- Lowerre, B. T. (1976). The Harpy Speech Recognition System. Ph.D. dissertation, Department of Computer Science, Carnegie Mellon University.
- Sutskever, I., Vinyals, O., & Le, Q. V. (2014). Sequence to sequence learning with neural networks. Advances in Neural Information Processing Systems (NeurIPS).
- Wu, Y., et al. (2016). Google's neural machine translation system: Bridging the gap between human and machine translation. arXiv preprint arXiv:1609.08144.