Decoding Temperature › Implementing Temperature-Controlled Sampling
Language AI HandbookAll chapters Decoding Temperature › Implementing Temperature-Controlled SamplingOne section · Sections in this chapter5 / 11 ↓
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
Let's build a complete temperature sampling implementation. We will start with a function that samples from temperature-scaled logits, then extend it to generate token sequences. The implementation is straightforward because temperature scaling is a single line of code: dividing the logits by before passing them to softmax.
Understanding the implementation details matters because they affect numerical stability. Very low temperatures can cause logits to become very large after scaling, potentially causing floating point overflow before exponentiation. Very high temperatures with very negative logits can cause the softmax denominator to approach zero. Professional implementations use the log-sum-exp trick to handle these edge cases, but for temperatures in the practical range (0.1 to 2.0), the naive implementation works well.
The log-sum-exp trick exploits the fact that for any constant . Choosing ensures all exponentiated values are between 0 and 1, preventing overflow. PyTorch's F.softmax applies this trick internally, so the naive division followed by softmax is numerically safe in practice.
def sample_with_temperature(
logits: torch.Tensor, temperature: float = 1.0
) -> int:
"""Sample a token index from logits with temperature scaling.
Args:
logits: Raw model output scores, shape (vocab_size,)
temperature: Scaling parameter. 1.0 = unchanged, <1 = sharper, >1 = flatter
Returns:
Sampled token index
"""
if temperature <= 0:
# Greedy selection for T <= 0
return logits.argmax().item()
# Apply temperature scaling
scaled_logits = logits / temperature
probs = F.softmax(scaled_logits, dim=0)
# Sample from the distribution
return torch.multinomial(probs, num_samples=1).item()The function handles the edge case of by returning the argmax (greedy decoding). For positive temperatures, it scales logits, converts to probabilities, and samples using PyTorch's multinomial function. The torch.multinomial function performs a single draw from a categorical distribution defined by the probability vector, which is exactly what we want.
# Demonstrate sampling behavior at different temperatures
def sample_distribution(
logits: torch.Tensor, temp: float, n_samples: int = 1000
) -> dict:
"""Sample many times and return frequency distribution."""
counts = {}
for _ in range(n_samples):
idx = sample_with_temperature(logits, temp)
counts[idx] = counts.get(idx, 0) + 1
return {k: v / n_samples for k, v in counts.items()}Empirical sampling frequencies (1000 samples each): Token T=0.5 T=1.0 T=2.0 -------------------------------------------- sunny 0.518 0.290 0.179 cloudy 0.209 0.184 0.163 rainy 0.056 0.092 0.119 nice 0.045 0.092 0.089 terrible 0.004 0.018 0.046 cold 0.016 0.059 0.078 warm 0.033 0.085 0.092 hot 0.014 0.041 0.068 perfect 0.101 0.127 0.129 unpredictable 0.004 0.012 0.037
The empirical frequencies closely match the theoretical probabilities computed earlier. At , samples cluster heavily on "sunny" and "cloudy." At , we see meaningful representation from lower-probability tokens like "warm," "nice," and "cold." This alignment between theoretical probabilities and empirical frequencies is expected: with 1,000 samples, the law of large numbers ensures the empirical distribution converges close to the theoretical one.
Batch Sampling for Efficiency
In practice, we often want to generate multiple completions or compare outputs across temperatures. Looping over a single-sample function is inefficient because it makes many small GPU operations instead of one large one. Here is a vectorized implementation that generates all samples in a single call to torch.multinomial:
def batch_sample_with_temperature(
logits: torch.Tensor, temperature: float = 1.0, num_samples: int = 1
) -> torch.Tensor:
"""Sample multiple tokens efficiently.
Args:
logits: Raw scores, shape (vocab_size,) or (batch, vocab_size)
temperature: Scaling parameter
num_samples: Number of samples to draw
Returns:
Tensor of sampled indices
"""
if temperature <= 0:
if logits.dim() == 1:
return logits.argmax().unsqueeze(0).expand(num_samples)
return logits.argmax(dim=-1).unsqueeze(-1).expand(-1, num_samples)
scaled = logits / temperature
probs = F.softmax(scaled, dim=-1)
if logits.dim() == 1:
return torch.multinomial(
probs, num_samples=num_samples, replacement=True
)
return torch.multinomial(probs, num_samples=num_samples, replacement=True)5 samples at each temperature: T = 0.3: rainy, sunny, sunny, sunny, rainy T = 1.0: sunny, warm, cloudy, hot, perfect T = 2.0: perfect, cloudy, rainy, sunny, warm
At low temperature, we see repeated "sunny" selections. Higher temperatures introduce variety, sometimes surfacing less expected but still contextually reasonable tokens. The batch implementation is much more efficient because torch.multinomial with replacement=True draws all samples in a single vectorized operation rather than looping.