Constrained Decoding › Constrained Sampling Strategies
Language AI HandbookAll chapters Constrained Decoding › Constrained Sampling StrategiesOne section · Sections in this chapter8 / 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
Beyond masking invalid tokens, we can combine constraints with sampling strategies like temperature, top-k, and nucleus sampling.
Temperature and Constraints
Temperature scaling adjusts the "sharpness" of the distribution. With constraints, we apply temperature before masking to maintain proper probability ratios among valid tokens.
import numpy as np
def constrained_sample_with_temperature(
logits: np.ndarray, valid_token_ids: set[int], temperature: float = 1.0
) -> int:
"""
Sample with both temperature and constraints.
Args:
logits: Raw model logits
valid_token_ids: Set of valid token indices
temperature: Sampling temperature (>1 = more random, <1 = more deterministic)
Returns:
Sampled token ID
"""
# Apply temperature first
scaled_logits = logits / temperature
# Mask invalid tokens
masked_logits = np.full_like(scaled_logits, float("-inf"))
for tid in valid_token_ids:
masked_logits[tid] = scaled_logits[tid]
# Convert to probabilities
max_logit = np.max(masked_logits[masked_logits > float("-inf")])
exp_logits = np.exp(masked_logits - max_logit)
probs = exp_logits / exp_logits.sum()
return np.random.choice(len(probs), p=probs)Effect of temperature on constrained sampling: Valid tokens: a, b, c, d, e (X is excluded) Original logits: a=2.0, b=1.5, c=1.0, d=0.5, X=3.5(invalid), e=0.2 Temperature 0.5: a: ████████████████████████████████ 64.0% b: ██████████ 21.0% c: █████ 10.4% d: █ 2.5% e: █ 2.1% Temperature 1.0: a: ████████████████████ 40.0% b: █████████████ 27.2% c: ████████ 16.6% d: ████ 8.7% e: ███ 7.5% Temperature 2.0: a: ███████████████ 31.3% b: ███████████ 23.4% c: █████████ 18.8% d: ███████ 14.5% e: ██████ 12.0%
Lower temperatures make the sampler more deterministic within the valid set, while higher temperatures spread probability more evenly across valid options. The constraint always prevents invalid tokens regardless of temperature.
Combining Top-k with Constraints
Top-k sampling and constraints interact in an important way: we should apply constraints first, then take top-k among valid tokens.
import numpy as np
def constrained_top_k_sample(
logits: np.ndarray,
valid_token_ids: set[int],
k: int = 5,
temperature: float = 1.0,
) -> int:
"""
Top-k sampling within constraint-valid tokens.
Args:
logits: Raw model logits
valid_token_ids: Set of valid token indices
k: Number of top tokens to consider
temperature: Sampling temperature
Returns:
Sampled token ID
"""
# First, mask to valid tokens only
valid_logits = []
valid_ids_list = []
for tid in valid_token_ids:
valid_logits.append(logits[tid])
valid_ids_list.append(tid)
valid_logits = np.array(valid_logits)
# Apply temperature
valid_logits = valid_logits / temperature
# Get top-k among valid tokens
if len(valid_logits) > k:
top_k_indices = np.argsort(valid_logits)[-k:]
top_k_logits = valid_logits[top_k_indices]
top_k_ids = [valid_ids_list[i] for i in top_k_indices]
else:
top_k_logits = valid_logits
top_k_ids = valid_ids_list
# Convert to probabilities and sample
exp_logits = np.exp(top_k_logits - np.max(top_k_logits))
probs = exp_logits / exp_logits.sum()
selected_idx = np.random.choice(len(top_k_ids), p=probs)
return top_k_ids[selected_idx]Top-k within constrained tokens: Logits: run=2.5, walk=2.0, jump=1.8, skip=1.5, hop=1.0, crawl=0.5, roll=0.2 Invalid: FLY=5.0 (highest logit but excluded) With k=3, sampling only from top-3 valid tokens: run: 47.3% walk: 28.1% jump: 24.6% Only 'run', 'walk', 'jump' are sampled (top-3 among valid).