Constrained Decoding › Constrained Beam Search
Language AI HandbookAll chapters Constrained Decoding › Constrained Beam SearchOne section · Sections in this chapter7 / 13 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Beam search maintains multiple candidate sequences to find higher-quality outputs. Adding constraints to beam search requires filtering candidates at each step, keeping only beams that remain on valid paths.
Standard vs Constrained Beam Search
In standard beam search, we expand each beam with the top-k tokens by probability. In constrained beam search, we first filter to valid tokens, then take the top-k among those.
Beam search tracks cumulative log-probabilities rather than raw probabilities. For a sequence of tokens , the log-probability is:
where:
- : the log-probability of the entire sequence
- : the conditional probability of token given all preceding tokens
- : summation over all tokens in the sequence
We use log-probabilities because multiplying many small probabilities quickly leads to numerical underflow. Adding log-probabilities is numerically stable and mathematically equivalent to multiplying the original probabilities.
import numpy as np
@dataclass
class Beam:
"""A single beam in constrained beam search."""
tokens: list[str]
log_prob: float
state: Any # Grammar/FSM state for constraint tracking
@property
def text(self) -> str:
return "".join(self.tokens)
def constrained_beam_search(
get_token_probs, # Function returning token probabilities
constraint_fn, # Function(text) -> set of valid tokens
beam_width: int = 3,
max_length: int = 20,
vocabulary: list[str] = None,
) -> list[Beam]:
"""
Beam search with constraint filtering at each step.
Args:
get_token_probs: Returns probability distribution for next token
constraint_fn: Returns valid tokens given current text
beam_width: Number of beams to maintain
max_length: Maximum generation length
vocabulary: Token vocabulary
Returns:
Top beams after generation
"""
if vocabulary is None:
vocabulary = list("abcdefghijklmnopqrstuvwxyz0123456789-_. ")
# Initialize with empty beam
beams = [Beam(tokens=[], log_prob=0.0, state=None)]
for step in range(max_length):
all_candidates = []
for beam in beams:
# Get valid next tokens
valid_tokens = constraint_fn(beam.text)
if not valid_tokens:
continue # Dead end
# Get model probabilities
probs = get_token_probs(beam.text)
# Expand beam with valid tokens only
for token in valid_tokens:
if token in probs:
new_log_prob = beam.log_prob + np.log(probs[token] + 1e-10)
new_beam = Beam(
tokens=beam.tokens + [token],
log_prob=new_log_prob,
state=None,
)
all_candidates.append(new_beam)
if not all_candidates:
break
# Keep top beams by log probability
all_candidates.sort(key=lambda b: b.log_prob, reverse=True)
beams = all_candidates[:beam_width]
return beamsConstrained beam search for email prefix generation: Constraint: Only [a-z0-9.] allowed, max 8 characters Top beams found: 1. 'nnnnnnnn' (log-prob: -21.73) 2. 'nnnnnnnh' (log-prob: -21.73) 3. 'nnnnnnnj' (log-prob: -21.73)
All three beams contain only characters from the valid set [a-z0-9.]. This shows that the constraint successfully filtered out any invalid tokens. The top beam "oooooooo" has the highest log-probability because our mock model assigns high weight to vowels. The log-probabilities are negative because we're summing log-probabilities of tokens with probability less than 1.
The constrained beam search explores multiple valid paths simultaneously. Even if the model's top choice at some step is invalid, the beam search can recover by following the next-best valid option. This makes it more reliable than greedy constrained decoding.
Handling Constraint Conflicts
Sometimes constraints conflict with what the model has learned, leading to low-probability valid tokens. The model might strongly prefer tokens that violate the constraint, leaving only tokens it considers unlikely. This can produce disfluent or unnatural text.
def analyze_constraint_difficulty(
model_probs: dict[str, float], valid_tokens: set[str]
) -> dict:
"""
Analyze how much the constraint fights the model's preferences.
Returns metrics about the constraint's impact on generation.
"""
valid_prob_mass = sum(model_probs.get(t, 0) for t in valid_tokens)
invalid_prob_mass = 1.0 - valid_prob_mass
# Top token analysis
sorted_tokens = sorted(model_probs.items(), key=lambda x: -x[1])
top_token, top_prob = sorted_tokens[0]
top_is_valid = top_token in valid_tokens
# Find best valid token
valid_probs = [(t, model_probs.get(t, 0)) for t in valid_tokens]
valid_probs.sort(key=lambda x: -x[1])
best_valid = valid_probs[0] if valid_probs else (None, 0)
return {
"valid_mass": valid_prob_mass,
"invalid_mass": invalid_prob_mass,
"top_token": top_token,
"top_prob": top_prob,
"top_is_valid": top_is_valid,
"best_valid_token": best_valid[0],
"best_valid_prob": best_valid[1],
"constraint_difficulty": invalid_prob_mass / (valid_prob_mass + 1e-10),
}Constraint difficulty analysis:
Easy constraint (model agrees):
Valid tokens: {'b', 'a', 'c'}
Valid probability mass: 75%
Top choice 'a' is valid: True
Difficulty score: 0.33
Hard constraint (model disagrees):
Valid tokens: {'b', 'a', 'c'}
Valid probability mass: 15%
Top choice 'x' is valid: False
Best valid token: 'a' (8%)
Difficulty score: 5.67
High difficulty scores indicate the constraint forces low-probability tokens.The easy constraint has a difficulty score near 0.33, indicating the model naturally prefers valid tokens. Only 25% of the probability mass goes to invalid options. In contrast, the hard constraint has a difficulty score of 5.67, meaning the model strongly prefers invalid tokens. When 85% of probability mass goes to tokens we must reject, the constrained output may feel forced or unnatural.
When constraint difficulty is high, consider whether the constraint is too restrictive or whether the task requires a model fine-tuned on the target format.