Sampling — Module 11.1

How This Lesson Fits the Module & Volume

After softmax, the model offers P(token | context). Sampling draws one token from that categorical distribution instead of always taking the argmax.

This lecture is the base decoder; temperature, top-k, and top-p are filters that reshape the distribution before you sample. Beam search is the main non-sampling alternative for search-based decoding.

Learning Objectives

By the end of this lesson, students should be able to:

Contrast greedy decoding with multinomial sampling.
Implement one generation step: logits → probs → sample.
Explain why pure sampling can produce rare nonsense tokens.
Relate sampling to the autoregressive inference loop.
Use seeds / generators for reproducible draws in PyTorch.
Preview how truncation strategies improve sample quality.

Definition

Sampling (ancestral / multinomial sampling) selects the next token by drawing from the categorical distribution P(· | x_{<t}) produced by the LM. If p = softmax(z), then token id ~ Categorical(p). Repeating this step builds a generated sequence.

Greedy vs Sample

Context

Prompt tokens

Forward

Logits for step t

Decide

Argmax or sample

Append

Grow the sequence

Strategy	Rule	Typical feel
Greedy	argmax p	Safe, often bland / repetitive
Pure sampling	draw from full p	Diverse, can be wild / low-quality
Tempered sampling	softmax(z/T) then draw	Controllable randomness
Truncated (top-k / top-p)	zero out tail, renormalize, draw	Production default for chat LMs

When to Sample

Prefer sampling

Creative writing, brainstorming.
Need multiple distinct completions.
Chat assistants (with truncation).

Prefer greedy / beam

Short factual answers.
Structured formats (JSON keys).
Classic MT / ASR pipelines (beam).

Hybrid practice

Sample with top-p + mild temperature.
Greedy for tool-call arguments.
Tune per product surface.

Code: One Sampling Step

import torch
import torch.nn.functional as F

def sample_next(logits, generator=None):
    # logits: (B, V)
    probs = F.softmax(logits, dim=-1)
    # multinomial expects probs; draws 1 token per batch row
    return torch.multinomial(probs, num_samples=1, generator=generator)

torch.manual_seed(0)
logits = torch.tensor([[3.0, 1.0, 0.5, -1.0]])  # (1, 4)
ids = [sample_next(logits).item() for _ in range(5)]
print(ids)  # varied draws, biased toward index 0

# Greedy contrast
greedy = logits.argmax(-1).item()
print("greedy:", greedy)

Strengths and Tradeoffs

Strengths

Matches the probabilistic LM training objective.
Natural diversity across runs.
Simple building block for all truncated samplers.

Tradeoffs

Long-tail tokens can derail quality.
Non-deterministic without fixed seeds.
Does not optimize whole-sequence score (unlike beam).

Common Misconception

“Sampling means the model is guessing randomly with equal odds.” Draws follow the model’s probabilities: high-mass tokens are far more likely. Randomness is weighted, not uniform—unless the distribution itself is nearly flat (e.g., extreme temperature).

Related module pages: Softmax, Temperature, Top-K, Top-P, Inference.

Knowledge Check

Short Answer: What distribution do we sample from? Answer: The categorical softmax distribution over the vocabulary.
True/False: Greedy decoding is a form of multinomial sampling. Answer: False—greedy is argmax, not a random draw.
Multiple Choice: torch.multinomial needs: (a) probabilities, (b) raw token strings, (c) learning rates. Answer: (a).
Short Answer: Name one risk of sampling from the full vocabulary. Answer: Drawing rare / low-quality tail tokens.
True/False: Sampling is weighted by the model’s probabilities. Answer: True.
Multiple Choice: Production chat LMs usually: (a) pure full-vocab sample, (b) truncated sample (top-k/p), (c) never decode. Answer: (b).
Short Answer: What happens after you sample token t? Answer: Append it to the context and run another forward step.
True/False: Fixing a RNG seed can make sampling reproducible. Answer: True.
Multiple Choice: Beam search primarily: (a) samples randomly, (b) keeps top partial sequences, (c) trains embeddings. Answer: (b).
Short Answer: Which lecture reshapes logits before softmax for sharper/flatter draws? Answer: Temperature.

Key Takeaways

Sampling draws the next token from softmax probabilities.
It enables diversity; greedy is the deterministic alternative.
Pure full-vocab sampling often needs truncation in practice.
Generation = repeated sample-and-append.
Next: Temperature.

Trainer’s Guide

Hands-on idea: Sample 20 tokens from a peaked vs nearly flat distribution; histogram the IDs and discuss coverage of the tail.

Discussion prompt: For a customer-support bot, when is non-determinism a product bug rather than a feature?

Recap: Sampling turns the next-token distribution into a concrete token ID. Continue with Temperature.