AI Engineering Foundations 16: Sampling — From a Uniform Die to GPT's 50,000-Word Vocabulary
This is lesson 16 of the series working through AI Engineering from Scratch (58k GitHub stars, MIT license), a 120-minute build-style lesson with probability and Bayes as prerequisites.
Fifty Thousand Scores, Pick One
A language model outputs one score per vocabulary token — about fifty thousand numbers. Always picking the highest makes a复读机-style repeater; pure randomness produces nonsense. The middle ground is defined by sampling, and this lesson rebuilds the generative dice from scratch.
One Uniform Number Is Enough for Any Distribution
A uniform random number on [0, 1) has mean 1/2 and variance 1/12 — humble but powerful: one uniform sample suffices to produce a sample from *any* target distribution, via inverse CDF sampling. The proof is three lines:
P(F⁻¹(U) ≤ x) = P(U ≤ F(x)) = F(x)
Practice with the exponential distribution: F(x) = 1 − e^(−λx), so x = −ln(u)/λ. Running 10,000 samples gave an empirical mean of 0.4995 vs. the theoretical 1/λ = 0.5. The catch: the normal distribution's CDF has no closed-form inverse, which motivates techniques like Box-Muller.
Propose, Then Roll the Dice
When you can't invert the CDF but can evaluate the target density up to a constant, rejection sampling applies: propose from a tractable q, envelope with M·q, accept when u < p(x)/(M·q(x)).
For the course's Beta(2, 5) exercise: peak density 2.4576 at x=0.2, so with a Uniform(0,1) proposal, M = 2.4576 and the theoretical acceptance rate is 1/M = 0.4069. An actual run of 120,000 proposals gave 0.4059 — theory and practice differ by one part in a thousand.
Two caveats worth flagging:
- The 1/M figure assumes a normalized target density. The course text says the target may be off by a constant factor — in that case the acceptance rate is that constant divided by M, not 1/M.
- Rejection sampling's acceptance rate decays exponentially with dimension — in high dimensions nearly all proposal volume gets discarded. Fine in 1–2 dimensions, hopeless past ten.
- Temperature: divide logits by T before softmax. Example with logits 2 and 1: T=1 gives probability ratio e ≈ 2.72; T=0.5 gives e² ≈ 7.39; T=2 gives √e ≈ 1.65. T < 1 sharpens; T > 1 flattens.
- Top-k: keep only the k highest-probability tokens, renormalize. Flexible in spirit but rigid in practice: k=40 wastes slots when the model is confident (95% on one token) and amputates valid options when probability spreads over a thousand tokens. k=50 was popularized by Fan et al.'s 2018 hierarchical story generation paper.
- Top-p (nucleus sampling): keep the smallest set of tokens whose cumulative probability exceeds p (e.g., 0.9). The candidate set grows or shrinks with the model's confidence. From Holtzman et al., 2020, *The Curious Case of Neural Text Degeneration*. Temperature 0.7 + top-p 0.9 is the most common general-purpose pairing.
Sometimes You Only Need Expectations
Importance sampling estimates expectations under p using samples from another distribution q, weighted by p(x)/q(x). This is the heart of PPO: trajectories come from the old policy, but you optimize the new one, so each sample is multiplied by the ratio of new-policy to old-policy probabilities, with clipping to keep the new policy from drifting too far. The cost hides in the variance: the further q is from p, the more a few huge weights dominate.
Ten Thousand Dice Rolls to Squeeze Out Pi
Monte Carlo at its plainest: scatter points in a unit square; the fraction inside the quarter circle times 4 estimates π (100,000 points gave 3.1502). For the ∫sin exercise (true value 2), repeating each sample count 200 times gave estimator standard deviations of 0.0101 at N=10,000 and 0.0032 at N=100,000 — ratio 3.12 vs. the theoretical √10 ≈ 3.16. Error scales as 1/√N: 100× more samples buys only 10× less error.
The claim that this rate is "dimension-independent" deserves care: it's the *convergence rate* that is dimension-independent; the hidden constant Var(f) can blow up with dimension.
Walking Along the Distribution
Metropolis-Hastings stops trying to sample directly and instead walks: from any point, tentatively move; accept moves with probability min(1, density ratio). Detailed balance — equal flow between any x and x' — guarantees the footprints converge to the target distribution.
Running it on a bimodal target (two equally weighted Gaussians): after 50,000 steps, 46.6% of samples landed in the right mode vs. the theoretical 50%. The imbalance comes from small step sizes struggling to cross the low-density valley — the proposal scale dilemma in action: too small crawls, too large gets stuck rejecting everything.
The course cites the optimal acceptance rate of ~0.234 for high-dimensional Gaussian targets. Empirically: 20 dimensions with scale 0.5 gave 0.278; 100 dimensions with scale 0.25 gave 0.212 — sliding toward 0.234 as dimension grows. The number comes from Roberts, Gelman & Gilks (1994), but rests on assumptions: independent Gaussian target, random-walk proposal, dimension → ∞.
Gibbs sampling, a relative on this road, updates one variable at a time from its conditional distribution — acceptance is always 100% — but gets stuck moving along coordinate axes when variables are strongly correlated.
The Three LLM Knobs
Language models emit raw logits before softmax.
z / temperature divides by zero), while the full code/sampling.py in the repo correctly guards T ≤ 0 by returning the argmax. The repo version is the good one.Moving Randomness Out of the Computation Graph
Sampling is non-differentiable, which would sever gradients in VAEs. The reparameterization trick rewrites z ~ N(μ, σ²) as z = μ + σ·ε with ε ~ N(0,1), making z an ordinary function of parameters with ∂z/∂μ = 1 and ∂z/∂σ = ε. Verified numerically against the analytical derivatives — this single move makes VAEs trainable.
For discrete categories, the Gumbel-Max trick adds Gumbel noise to log-probabilities and takes the argmax; the landing frequencies exactly equal the original probabilities. With [0.6, 0.3, 0.1] over 50,000 draws: 0.604 / 0.299 / 0.098.
Another real bug on the course page: the Gumbel-Softmax example takes math.log(p) on inputs already named logits — a double log. Feeding it the course's own example logits [2, 1, 0.5, 0.1, −1] crashes on log(−1); with all-positive [3, 1, 0.1], frequencies over 5,000 draws were 0.73 / 0.24 / 0.03 vs. the true 0.84 / 0.11 / 0.05 — the second place inflated ~2×. The repo's full version correctly takes log_probs directly. The fix is simple: add Gumbel noise straight to logits and take the max — the softmax log-denominator is a constant and doesn't affect ordering. Gumbel-Softmax is the differentiable version: softmax with temperature τ in place of argmax, τ → 0 approximating one-hot, with straight-through estimation for hard forward / soft backward.
Partitioning Space for Forced Uniformity
Standard Monte Carlo covers space by luck. Stratified sampling splits [0,1] into N strata and samples one point from each. For the sin integral with 100 samples over 300 repetitions: standard estimator SD 0.091 vs. stratified 0.0020 — over 2000× variance reduction. NeRF uses it for sampling along rays.
Diffusion Models: Every Step Is Sampling
Zooming out: diffusion models generate images by a forward process adding Gaussian noise and a learned reverse process denoising step by step. Every denoising step uses reparameterization (one noise sample, one deterministic transform), the noise schedule plays temperature annealing, the training objective is a Monte Carlo-estimated evidence lower bound, and the whole generation chain is a Markov chain. Image generation is iterative sampling end to end.
A Parting Exercise
Exercise 4: sample p(x,y) ∝ exp(−(x²y² + x² + y² − 8x − 8y)/2) with Metropolis-Hastings. Direct sampling is hopeless for this lopsided distribution — but conditionals on x (given y) and y (given x) are exactly Gaussian, so Gibbs sampling solves it precisely. One problem, two roads, each with its own temperament.
Fifty thousand word cards; only a handful of dice. See you next lesson.