--- title: Stochastic Beam Search (SBS) url: https://www.emergentmind.com/topics/stochastic-beam-search-sbs type: topic --- # Stochastic Beam Search (SBS) Stochastic Beam Search (SBS) is a family of decoding algorithms for sequence models that blend the structured exploration of beam search with the stochasticity of sampling. The method addresses longstanding limitations of deterministic beam search—such as low diversity and overemphasis on a narrow mode of the model’s distribution—while avoiding the inefficiency and high variance of pure sampling. SBS has foundational connections to the Gumbel-Top-k trick for sampling without replacement, and recent variants further incorporate evaluation signals (e.g., LLM self-critique) to robustly guide multi-step reasoning, especially in large language models. The SBS framework unifies a spectrum of inference regimes, including deterministic beam search as a limiting case, and is now a critical component in settings requiring diverse candidate generation, low-variance expectation estimation, and calibrated multi-step reasoning. ## 1. Core Principles and Algorithmic Foundations Stochastic Beam Search operates by propagating a set (beam) of $k$ partial hypotheses through the exponentially large space of possible sequences. At each decoding step, a larger candidate set is generated—typically by expanding each beam item with multiple possible continuations—and a stochastic, *without-replacement* selection is performed using a probability distribution over these candidates. The distribution is frequently parameterized by a temperature hyperparameter to interpolate between deterministic top-k (pure exploitation) and uniform random selection (pure exploration) [1903.06059, 2305.00633]. A pivotal technique underlying SBS is the *Gumbel-Top-k* trick, which enables sampling $k$ items without replacement from a categorical or structured sequence distribution. At each expansion, Gumbel noise is added to (potentially unnormalized) log-probabilities, and the $k$ highest perturbed scores determine the selected candidates. This approach preserves the marginal sequence distribution while enforcing sample diversity [1903.06059]. Variants such as Conditional Poisson Stochastic Beam Search generalize the candidate sampling operation by employing conditional Poisson designs, giving rise to rigorously specified inclusion probabilities for unbiased downstream estimation [2109.11034]. ## 2. Mathematical Formulation In standard sequence modeling, the probability of a sequence $R = [s^1, \ldots, s^T]$ given prompt $x$ is $$ P(R | x) = \prod_{t=1}^T P(s^t | x, s^{1:t-1}) $$ SBS defines a *stochastic pruning* operator at each step. At decoding step $t$, for each of the current $k$ beams, $n$ continuations are sampled, forming $kn$ candidates. Each candidate $i$ is assigned an accumulated log-score $L(R^{1:t}_i)$ (optionally incorporating auxiliary signals such as self-evaluation confidence). The normalized selection probability for candidate $i$ is $$ P_{\text{beam}}(R^{1:t}_i) = \frac{\exp(L(R^{1:t}_i) / \tau)}{\sum_{j \in S} \exp(L(R^{1:t}_j) / \tau)} $$ where $\tau$ is a temperature parameter. $k$ sequences are sampled without replacement from this distribution to form the next beam [2305.00633]. In Gumbel-Top-k SBS, each sequence candidate is perturbed by an independent Gumbel(0,1) variable and the $k$ highest perturbed scores are selected. For sequence models, this is implemented implicitly in a top-down manner, maintaining only $O(kL)$ model evaluations, where $L$ is the sequence length [1903.06059]. Conditional Poisson SBS replaces the softmax/Gumbel sampling with an exact conditional Poisson sampling scheme, robustly enabling unbiased estimation of expectations and more stable sample inclusion probabilities [2109.11034]. ## 3. Algorithmic Procedure The typical SBS decoding cycle is as follows (parameters: $k$ beam size, $n$ rollouts per beam, $\tau$ temperature, $\lambda$ optional mixture weight): - **Initialization**: Start with an empty sequence or prompt. Set initial $\tau$ (e.g., $0.2 \leq \tau_0 \leq 0.5$). - **Expansion**: For each beam, sample $n$ candidate continuations from the sequence model. - **Scoring**: For each candidate, compute the accumulated log-probability, optionally augmenting with external guidance. In self-evaluation-guided SBS, the score is a mixture of model likelihood and a correctness confidence produced by a verifier model: $$ L(R^{1:t}) = \sum_{j=1}^t \left[ \lambda \log P(s^j | x, s^{1:j-1}) + (1-\lambda) \log C(s^j) \right] $$ - **Sampling**: Compute a temperature-controlled softmax over all candidate log-scores. Sample $k$ distinct items without replacement according to these probabilities. - **Annealing**: Decrease $\tau$ (e.g., $\tau \leftarrow \alpha \tau$ with $0 < \alpha < 1$), shifting gradually from exploration to exploitation. - **Iteration**: Repeat for each sequence step up to a preset maximum, then collect final beam. Key resource consumption per step is $O(k n)$ forward passes and scoring operations. The overall cost is $O(T k n)$ in model invocations, where $T$ is the sequence depth [2305.00633]. ## 4. Temperature Hyperparameter and Exploration–Exploitation Balance The temperature $\tau$ directly governs the exploration–exploitation trade-off in candidate selection: - As $\tau \to 0$, the softmax curve sharpens, recovering deterministic beam search where the $k$ highest-scoring candidates are chosen. - As $\tau$ increases, lower-scoring candidates gain weight, increasing diversity and the likelihood of escaping local maxima. - High initial $\tau$ with subsequent annealing (e.g., $\alpha \approx 0.5$ per step) balances early diversity with later-stage focus, empirically improving multi-step reasoning and mitigating error accumulation [2305.00633]. This mechanism parallels the role of temperature in softmax sampling and stochastic optimization, and can be adapted for different diversity/accuracy requirements. ## 5. Variants and Theoretical Connections ### Gumbel-Top-k Stochastic Beam Search - Employs the Gumbel-Top-k trick to sample $k$ sequences without replacement from posited sequence probabilities. - Bridges the regime between deterministic beam search and pure sequence sampling, interpolating from mode-seeking to distribution-covering behavior as $\tau$ is varied [1903.06059]. - Enables low-variance, unbiased estimators for expectations of sequence-level functions (e.g., expected BLEU, entropy) by leveraging the properties of without-replacement sampling. - Theoretical cost is linear in $k$ and $L$ (sequence length), avoiding exponential blowup [1903.06059]. ### Conditional Poisson Stochastic Beam Search - Replaces the softmax or Gumbel procedure with a conditional Poisson sampling process, allowing for exact computation of inclusion probabilities and unbiased Horvitz–Thompson estimation. - Reduces variance in expectation estimation tasks, especially under high entropy, compared with Gumbel-based SBS and ancestral sampling. - Adds $O(NK)$ dynamic programming overhead, but remains practical with vocabulary truncation [2109.11034]. #### Comparative Table | Decoding Variant | Sampling Scheme | Variance in Expectation Estimators | |------------------------------|------------------------|------------------------------------| | Deterministic Beam Search | Top-k deterministic | Biased, low coverage | | SBS (Gumbel-Top-k) | Gumbel w/o replacement | Lower variance, some bias | | Conditional Poisson SBS | Cond. Poisson w/o repl.| Lowest variance, unbiased | ## 6. Role of Self-Evaluation and Guided Reasoning Recent adaptations of SBS integrate stepwise *self-evaluation guidance* for multi-step reasoning, notably in large language models. At each step: - The continuation $s^t$ is both sampled and scored by the base model and evaluated by a self-evaluation LLM or component, which outputs a correctness confidence $C(s^t) \in [0,1]$. - The score used in pruning and sampling is a convex combination of model likelihood and verifier confidence, calibrated via parameter $\lambda$: $$ E(R^{1:T}) = \prod_{t=1}^T \left[ P(s^t | x, s^{1:t-1})^{\lambda} \cdot C(s^t)^{1-\lambda} \right] $$ - This mechanism inhibits the propagation of early-step reasoning errors and increases robustness, particularly evident in benchmarks for arithmetic, symbolic, and commonsense reasoning. Improvements over Codex-backboned baselines include accuracy increases of 6.34%, 9.56%, and 5.46% on GSM8K, AQuA, and StrategyQA, respectively [2305.00633]. Self-evaluation-guided SBS realizes a lightweight, model-agnostic calibration strategy that integrates fluency and logical soundness without requiring model finetuning. ## 7. Applications, Trade-offs, and Empirical Insights SBS is broadly applied in: - Neural machine translation, where it outperforms deterministic and diverse beam search in balancing BLEU and output diversity, and provides lower-variance estimators for sentence-level metrics [1903.06059, 2109.11034]. - Multi-step reasoning in LLMs, where it reduces error accumulation and increases solution consistency [2305.00633]. - Settings requiring diverse solution sets, such as minimum Bayes risk decoding and expectation estimation in RL. Key trade-offs include: - Larger $k$ and $n$ yield better coverage and accuracy but increase computational cost and latency. - Smaller $\tau$ expedites convergence but risks omitting viable alternatives. - Self-evaluation injects additional forward passes for confidence computation. Empirical results indicate that SBS (with or without enhancements like conditional Poisson sampling or self-evaluation) achieves superior diversity–quality profiles and substantially lower estimator variance than classical alternatives [1903.06059, 2109.11034, 2305.00633]. --- **References:** - "Self-Evaluation Guided Beam Search for Reasoning" [2305.00633] - "Conditional Poisson Stochastic Beam Search" [2109.11034] - "Stochastic Beams and Where to Find Them: The Gumbel-Top-k Trick for Sampling Sequences Without Replacement" [1903.06059] Source: https://www.emergentmind.com/topics/stochastic-beam-search-sbs