---
title: Temperature-Annealed Stochastic Beam Search
url: https://www.emergentmind.com/topics/temperature-annealed-stochastic-beam-search-sbs
type: topic
---

# Temperature-Annealed Stochastic Beam Search

Temperature-Annealed Stochastic Beam Search (SBS) extends the classical beam search paradigm by integrating the Gumbel-Top-$k$ trick and temperature annealing, enabling sampling of $k$ distinct sequences without replacement and facilitating fine-grained exploration–exploitation trade-offs. This approach generalizes beam search into a stochastic sampling regime and provides a principled method for generating diverse, high-quality sequence outputs in applications spanning sequence modeling and protein engineering [1903.06059, 2603.10302].

## 1. Mathematical Foundations: Gumbel-Top-$k$ and Stochastic Sampling

The core of temperature-annealed SBS is the Gumbel-Max and Gumbel-Top-$k$ tricks. In Gumbel-Max sampling, un-normalized log-probabilities $\varphi_i$ for categories $i=1,\dots,n$ are perturbed by independent Gumbel(0) noise $G_i$, yielding samples from the categorical distribution via
$$
i^* = \arg\max_i\, \{\varphi_i + G_i\}.
$$
For sampling $k$ elements without replacement, the Gumbel-Top-$k$ trick takes the $k$ indices with the largest perturbed scores:
$$
(i_1^*, ..., i_k^*) = \text{argtop}_k\, \{\varphi_i + G_i\},
$$
which produces an ordered sample without replacement, matching the product of conditional probabilities over without-replacement pools. This foundation enables direct stochastic analogues of deterministic ranking-based search procedures [1903.06059].

For sequence models, where the distribution is factorized as $p_\theta(y_{1:T}) = \prod_{t=1}^T p_\theta(y_t | y_{1:t-1})$, the Gumbel process can be extended using a recursive tree construction. For a partial sequence/hypothesis $S$, the randomized score can be written as
$$
\varphi_S = \log \sum_{i \in S} \exp \varphi_i,
$$
with $G_{\varphi_S} \sim \mathrm{Gumbel}(\varphi_S)$ defining i.i.d. randomized scores for partial beams.

## 2. Algorithmic Structure and Pseudocode

Stochastic Beam Search (SBS) uses the above Gumbel-based formulation to sample $k$ paths without replacement. At each step, candidate extensions are scored by their (possibly temperature-scaled) log-probabilities, with independent Gumbel noise injected for stochastic ranking. The procedure selects the top-$k$ candidates into the beam for each subsequent decoding step. This process is formally described for autoregressive models in [1903.06059] and for masked language models (MLMs) in [2603.10302].

In the context of masked language models for sequence design, e.g., protein engineering:
- Start from a seed sequence $x^{(0)}$.
- At each of $E$ beam steps, replace each beam member with all one-substitution neighbors, rapidly scoring them via approximate pseudo-log-likelihoods under the MLM at temperature $T_k$.
- For each candidate, perturb its score with independent Gumbel noise.
- Select the top $B$ sequences for the next beam.

This generalizes directly to scenarios with multiple optimization objectives by allowing the ranking criterion to integrate, for instance, physicochemical metrics alongside model PLL [2603.10302].

## 3. Temperature Scaling and Annealing Schedules

A central feature of temperature-annealed SBS is the explicit control of the softmax temperature $\tau$ (or $T$ for MLMs) at each search step. Temperature scaling modifies the sharpness of the underlying categorical distribution:
$$
p_\theta(y_t|y_{1:t-1}; \tau) = \frac{\exp[\varphi_\theta(y_t|y_{1:t-1})/\tau]}{\sum_{y'} \exp[\varphi_\theta(y'|y_{1:t-1})/\tau]}.
$$
Lower $\tau$ sharpens the distribution (favoring high-probability outputs, exploitation), while higher $\tau$ flattens it (promoting diverse exploration). In the SBS algorithm, this translates into dividing all logit scores by $\tau$ before adding Gumbel noise [1903.06059, 2603.10302].

Annealing schedules for the temperature permit gradual modulation from exploration to exploitation. Empirical practice and algorithm design commonly use:
- **Linear schedule**: $\tau(t) = \tau_{\text{start}} - (\tau_{\text{start}}-\tau_{\text{end}})\cdot(t-1)/(T_{\max}-1)$
- **Exponential schedule**: $\tau(t) = \tau_{\text{start}} \cdot \exp(-\alpha (t-1))$
- **Inverse-time**: $\tau(t) = \tau_{\text{start}}/(1 + \beta (t-1))$
with hyperparameters tuned for application-specific trade-offs [1903.06059, 2603.10302].

## 4. Integration with Masked Language Models and Efficient Neighborhood Evaluation

For MLM-based design and optimization, SBS exploits the ability to rapidly evaluate the approximate pseudo-log-likelihood (PLL) of the $1$-edit neighborhood of a given sequence. The PLL of candidate sequence $s'$ under temperature $T_k$ is computed as
$$
\mathrm{PLL}(s'; T_k) \approx \sum_{i=1}^L \log \mathrm{softmax}_{T_k}(z_i)[s'_i],
$$
where $z_i$ are the masked logits at position $i$ obtained from the template $s$. Utilizing the wild-type marginal approximation, MLM calls are amortized across neighbors, substantially reducing computation from $O(L^4)$ to $O(BL^3)$ per beam step [2603.10302].

This efficient batched evaluation facilitates simultaneous scoring of all possible one-substitution mutants, which are then perturbed and ranked as per the Gumbel-Top-$k$ trick.

## 5. Hyperparameters, Trade-offs, and Computational Complexity

Key hyperparameters in temperature-annealed SBS are beam width $B$, maximum number of edits $E$, and the temperature schedule parameters ($T_0$, $T_E$, or $\alpha$ for exponential decay). Their tuning directly affects exploration depth, diversity, and computational cost:
- **Beam width ($B$)**: Larger $B$ increases diversity and search coverage but scales computational complexity; typical values are $B=5$–$20$.
- **Maximum edits ($E$)**: Controls search radius from the seed; higher $E$ increases off-distribution exploration but may yield unrealistic candidates.
- **Temperature schedule**: High initial $T_0$ supports diversity; final $T_E$ enforces convergence towards high-likelihood outputs.

An increase in $B$ or $T$ tends to produce more diverse outputs but may decrease the likelihood quality. Decreasing $B$ or the final $T_E$ emphasizes exploitation of the model's pseudo-likelihood objective. Computationally, SBS requires $O(BEL^3)$ operations for $E$ steps, orders of magnitude faster than mutation-centric sampling for generating diverse candidate sets [2603.10302].

Practical considerations include caching per-beam MLM logits, tuning $B$, $E$, and temperature schedules for experimental throughput, and integrating additional oracles or objectives via weighted scoring before the stochastic perturbation step.

## 6. Empirical Effects and Theoretical Properties

Temperature-annealed SBS exhibits nuanced empirical and theoretical properties:
- Increasing $\tau$ (or $T$) raises output diversity (e.g., unique $n$-grams in translation or more sequence diversity in protein design) but may reduce mean quality metrics such as BLEU or PLL.
- For moderate $\tau$ values ($\sim$0.2–0.5), SBS retains nearly optimal oracle performance while achieving substantially higher diversity than classical beam or Diverse Beam Search in language tasks [1903.06059].
- Large temperature settings can lower the variance of priority-sampling estimators for downstream metrics (BLEU, entropy), but sample distribution is biased toward lower-probability regions.
- Theoretically, as $\tau\to 0$, SBS converges to deterministic beam search; as $\tau \to \infty$, it approximates uniform sampling over the search space, with annealing providing a principled interpolation between these extremes [1903.06059].

In protein engineering, SBS with temperature tuning and efficient PLL approximation produces diverse and high-likelihood sequence libraries, outperforming mutation-centric sampling both in compute efficiency and quality-diversity trade-off [2603.10302].

## 7. Applications and Extensions

Temperature-annealed SBS has demonstrated practical impact in translation tasks, where it yields diverse, high-quality hypotheses, and in protein sequence engineering, where it efficiently generates variant libraries for in silico and in vitro selection [1903.06059, 2603.10302]. Its plug-in nature allows integration with complex scoring functions, supports multi-objective optimization, and is applicable wherever candidate sequences (from language to biological domains) are scored by autoregressive or masked models.

The method’s amenability to temperature control, low-variance estimators, and theoretical guarantees on without-replacement sampling position it as a robust intermediate between deterministic search and pure sampling, supporting both exploration and exploitation flexibly within large hypothesis spaces.

Source: https://www.emergentmind.com/topics/temperature-annealed-stochastic-beam-search-sbs