Beam Search

The beam width is the maximum number of partial sequences (hypotheses) retained at each decoding step. It is the algorithm's primary trade-off knob.

• Small k (e.g., 1): Equivalent to greedy search. Fast and memory-efficient but prone to suboptimal local maxima.
• Large k (e.g., 10, 50): Explores more of the search space, increasing the likelihood of finding a higher-probability global sequence. Exponentially increases memory and compute requirements. In robotics, a moderate beam width (e.g., 5-10) is often used to find robust action sequences without prohibitive latency.

Sequence probabilities naturally decrease with length, causing the algorithm to unfairly prefer shorter outputs. Length normalization counteracts this bias.

• The total score for a sequence of length T is typically adjusted: score = (1/T) * Σ log(P(token)) or score = (1/T^α) * Σ log(P(token)).
• The length penalty hyperparameter (α) controls the strength of normalization. Common values are between 0.6 and 1.0.
• Without normalization, beam search will prematurely terminate sequences, which is catastrophic for multi-step robotic action plans.

Managing when hypotheses are considered complete is crucial for variable-length output, like action plans.

• When a hypothesis generates an end-of-sequence (EOS) token, it is removed from the active beam and placed in a pool of completed sequences.
• The search can continue until all beams hit EOS or a maximum length is reached.
• A key decision is whether to allow beams of different lengths to compete directly or to apply length normalization for a fair comparison.

To reduce unnecessary computation, beam search often employs early stopping criteria.

• Global Early Stopping: The search halts when a specified number of complete sequences have been collected.
• Hypothesis Pruning: Low-probability beams can be discarded if their score falls below a threshold relative to the best beam.
• In time-sensitive robotic applications, a maximum decoding steps limit is strictly enforced to guarantee a response within a control cycle.

Standard beam search can lack diversity, with multiple beams converging on similar sequences. Diverse Beam Search introduces a diversity-promoting penalty.

• Beams are grouped. The score for a hypothesis in a group is penalized based on its similarity to higher-ranked hypotheses in previous groups.
• This is controlled by a diversity strength hyperparameter and the number of beam groups.
• For robotics, this can help discover qualitatively different action strategies for the same high-level instruction.

While not a parameter of beam search itself, the temperature of the model's output softmax directly affects the search landscape.

• Low Temperature: Sharpens the probability distribution, making scores between the top-k beams more distinct. This can make beam search more decisive but also more brittle.
• High Temperature: Flattens the distribution, giving more beams similar probabilities. This increases exploration but can lead to noisier, lower-quality final sequences.
• The temperature is typically tuned in conjunction with beam width for optimal results.

Beam search is the standard decoding algorithm for autoregressive language models (e.g., GPT, T5) and machine translation systems. It balances the exhaustive search of greedy decoding with the intractable search space of exploring all possibilities.

• Core Function: At each generation step, it maintains the top-k (beam width) most probable partial sequences.
• Trade-off: A larger beam width increases the chance of finding a higher-probability output but increases computational cost and can lead to repetitive or generic text.
• Example: Used by models like BERT for masked token prediction and in sequence-to-sequence architectures for summarization and dialogue.

In robotics, beam search is applied to decode sequences of action tokens from vision-language-action (VLA) models. It finds the most probable multi-step action plan given a high-level instruction.

• Process: The model generates a probability distribution over possible next action tokens (e.g., GRASP, MOVE_LEFT). Beam search explores multiple plausible future action sequences in parallel.
• Advantage Over Greedy: Prevents committing to a single, potentially sub-optimal initial action. Crucial for long-horizon tasks where early mistakes are catastrophic.
• Integration: Often used with Decision Transformers or autoregressive policy networks that output tokenized actions.

Beam search is fundamental in connectionist temporal classification (CTC) and sequence-to-sequence models for converting audio waveforms to text transcripts.

• Handles Alignment: Manages the alignment problem between variable-length audio frames and output characters/words.
• Context Integration: Modern systems use beam search with an external language model (shallow fusion) to bias decoding towards grammatically correct and contextually appropriate text.
• Scale: Deployed in real-time systems, requiring optimized beam search implementations to meet latency constraints.

When AI models like Codex or Code Llama generate code, beam search helps produce syntactically correct and logically consistent programs.

• Challenges: Code has strict syntax and long-range dependencies. Beam search evaluates multiple partial code completions.
• Combined with Constraints: Often paired with action masking to invalidate tokens that would lead to syntax errors (e.g., an unmatched parenthesis).
• Output Diversity: Using a moderate beam width can generate several valid coding alternatives for a single prompt.

Beam search is a core technique in heuristic search for games and puzzle solving, where the state space is too large for brute-force methods.

• Classic Use: Planning in board games or puzzle environments. It prunes the search tree to a manageable width.
• Modern Integration: Used within learned models, such as a policy network proposing candidate moves, which are then evaluated and expanded via beam search.
• Example: Algorithmic play in games like Sokoban or for generating strategic sequences in real-time strategy games.

The beam width (k) is the critical hyperparameter controlling the algorithm's behavior and computational cost.

• Small k (e.g., 1-5): Faster, less memory, but risk of search error (missing high-quality sequences). k=1 is greedy decoding.
• Large k (e.g., 10-100): Higher quality results, more stable, but with diminishing returns, increased latency, and memory usage.
• Practical Tuning: Chosen based on task criticality and latency budgets. Robotics often uses smaller beams (3-10) for real-time control, while offline translation may use larger beams.
• Advanced Variants: Diverse Beam Search modifies the scoring to promote variety among the k beams, preventing redundant sequences.

The fundamental sequential generation process where a model predicts the next token conditioned on all previously generated tokens. This is the core mechanism beam search operates within.

• Core Mechanism: The model's output distribution for step t+1 depends on the sequence of tokens generated from step 1 to t.
• Beam Search Context: Beam search manages multiple candidate sequences during this autoregressive process, expanding and pruning them at each step.
• Action Generation: In robotics, this translates to predicting the next action token (e.g., a joint angle delta) based on the history of previous actions and the current visual-language context.

A deterministic search strategy that selects the single most probable token at each step of autoregressive generation.

• Algorithm: At every generation step, the token with the highest predicted probability is chosen, forming a single sequence.
• Comparison to Beam Search: Greedy decoding is a special case of beam search with a beam width of 1. It is computationally cheap but can lead to suboptimal overall sequences, as early high-probability choices may force later low-probability ones.
• Use Case: Often used for fast, deterministic inference where optimality is less critical than speed.

A technique for controlling the randomness (entropy) of a model's output distribution during token sampling.

• Mechanism: The model's logits are divided by a temperature parameter T before applying the softmax function. T > 1 flattens the distribution (more diverse, creative outputs). T < 1 sharpens it (more deterministic, focused outputs). T → 0 approximates greedy decoding.
• Interaction with Beam Search: Temperature can be applied to the model's output distributions before beam search ranks and selects the top-k candidates. It directly influences which sequences are considered 'promising'.
• Application: Used to tune the trade-off between plan diversity and reliability in action sequence generation.

Stochastic decoding methods that sample from a restricted subset of the vocabulary to avoid low-probability tokens.

• Top-k Sampling: Filters the vocabulary to only the k tokens with highest probability at each step, then re-normalizes and samples from this set.
• Top-p (Nucleus) Sampling: Dynamically selects the smallest set of tokens whose cumulative probability exceeds p (e.g., 0.9), then samples from this set. Adapts to the shape of the distribution each step.
• Contrast with Beam Search: These are sampling-based methods that produce a single, diverse sequence. Beam search is a deterministic search over multiple sequences. They are often used as alternatives to beam search for more creative or varied text generation, but beam search is typically preferred for technical outputs like action plans where likelihood is paramount.

A technique that enforces physical or safety constraints by preventing a policy from selecting invalid actions during decoding.

• Mechanism: Before selecting an action (or token), the probabilities for invalid actions are set to zero (-inf in log space). This 'mask' is applied to the model's output distribution.
• Integration with Beam Search: The mask is applied at each step of beam expansion. This ensures every candidate sequence in the beam remains physically feasible, preventing the waste of computational resources on invalid paths.
• Example: Masking out joint angle tokens that would cause self-collision or exceed torque limits.

A transformer-based architecture that frames sequential decision-making as conditional sequence generation, a natural fit for beam search decoding.

• Paradigm Shift: Models the problem as Return-to-Go conditioned autoregressive prediction. It takes a desired return (reward), past states, and actions as input to predict future actions.
• Beam Search Application: When generating an action plan, beam search can be used over the Decision Transformer's autoregressive outputs to find the sequence of actions most likely to achieve the specified target return.
• Relevance to VLA Models: Provides a direct bridge from high-level language instructions ('achieve a high score on this task') to optimized action sequences, with beam search identifying the most probable successful trajectory.