Probability Distribution — Module 11.1

How This Lesson Fits the Module & Volume

The previous lecture defined a language model as a machine that outputs P(next token | context). That object is a categorical probability distribution over the vocabulary. Everything later in Module 11.1—softmax, temperature, top-k, top-p, and sampling—is about shaping or drawing from that distribution.

You already met softmax and cross-entropy in classification (Vol. 05–06). Here the “classes” are vocabulary entries, and there is one distribution per position along the sequence.

Learning Objectives

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

State the axioms a next-token distribution must satisfy (non-negative, sums to 1).
Relate raw logits to probabilities via softmax.
Interpret probability mass, mode, entropy, and long-tail vocabulary mass.
Compute and visualize a tiny categorical distribution in PyTorch.
Explain why argmax (greedy) and sampling behave differently on the same distribution.
Connect distributions to the next lecture’s next-token prediction objective.

Definition

A probability distribution over a discrete vocabulary V is a vector p ∈ R^|V| with p_i ≥ 0 for all i and ∑_i p_i = 1. In an LM, p_i = P(token = i | context) at a given position.

From Logits to Probabilities

The Transformer’s final linear layer produces logits—unnormalized scores. Softmax converts them into a valid distribution:

p_i = exp(z_i) / ∑_j exp(z_j)

Training pushes p toward the one-hot (or soft) target for the true next token using cross-entropy. At inference you either take argmax p or sample from p (optionally after reshaping it).

Object	Shape (batch=1)	Meaning
Logits `z`	`(\|V\|,)`	Raw scores; can be negative
Probabilities `p`	`(\|V\|,)`	Non-negative; sum to 1
Mode	scalar id	`argmax p`—greedy choice
Entropy `H(p)`	scalar	Uncertainty; high = flatter

Peakiness vs Spread

Peaked (low entropy)

One token dominates (e.g. 0.92).
Greedy and sampling often agree.
Common after strong context cues.

Flat (high entropy)

Many tokens share mass.
Sampling explores; greedy can feel dull.
Typical mid-sentence open choices.

Long tail

Tiny mass over rare IDs.
Can produce weird samples if unfiltered.
Motivation for top-k / top-p later.

Code: Softmax, Entropy, and a Draw

import torch

# Toy vocab: 0=the, 1=a, 2=cat, 3=dog, 4=runs
logits = torch.tensor([2.5, 1.0, 0.2, 0.1, -1.0])
probs = torch.softmax(logits, dim=-1)
entropy = -(probs * probs.clamp_min(1e-12).log()).sum()

print("probs:", probs.round(decimals=4).tolist())
print("sum:", float(probs.sum()))
print("mode id:", int(probs.argmax()), "p=", float(probs.max()))
print("entropy:", float(entropy))

# One sample from the categorical distribution
sample_id = torch.multinomial(probs, num_samples=1).item()
print("sample:", sample_id)

Joint Sequence Probability

An autoregressive LM defines a joint distribution over full strings by chaining conditionals. The log probability of a completed sequence is the sum of log next-token probs—exactly what teacher-forced training maximizes in expectation. This is why perplexity (exp of average negative log-likelihood) is a natural LM metric.

Why Distributions Help

Uncertainty is explicit (entropy, top-p mass).
Composable with search (beam) and sampling.
Calibrated scores enable ranking and filtering.

Pitfalls

Mis-calibrated models: high p ≠ correct fact.
Numerical issues: use log-softmax in losses.
Huge |V| makes naïve visualization hard.

Common Misconception

“The model outputs probabilities directly.” Networks output logits. Probabilities appear only after softmax (or an equivalent normalization). Mixing up the two leads to broken temperature scaling and incorrect training losses—always apply cross-entropy to logits with F.cross_entropy / log_softmax, not to already-softmaxed values twice.

Knowledge Check

Short Answer: State the two requirements for a discrete probability vector. Answer: Entries ≥ 0 and entries sum to 1.
True/False: Logits must already sum to 1. Answer: False.
Multiple Choice: Softmax turns logits into: (a) embeddings, (b) a categorical distribution, (c) attention masks. Answer: (b).
Short Answer: What is the mode of a distribution? Answer: The outcome with maximum probability (argmax).
True/False: Higher entropy means a peakier next-token distribution. Answer: False—higher entropy is flatter / more uncertain.
Multiple Choice: torch.multinomial(probs, 1): (a) returns argmax, (b) draws a sample from probs, (c) sorts the vocab. Answer: (b).
Short Answer: How does an LM get a joint probability for a full string? Answer: Product (or sum in log space) of successive next-token probabilities.
Short Answer: Why prefer F.cross_entropy(logits, target) over softmax-then-NLL? Answer: Numerical stability via fused log-softmax; avoids underflow / double softmax bugs.
Multiple Choice: Long-tail vocabulary mass motivates: (a) larger batch size, (b) top-k / top-p truncation, (c) removing LayerNorm. Answer: (b).
True/False: A valid probability distribution can contain a negative entry. Answer: False.

Key Takeaways

LM outputs are categorical distributions over the vocabulary at each position.
Softmax maps logits → probabilities; training uses cross-entropy on those conditionals.
Entropy and tail mass explain why greedy vs sampling (and later top-k/p) matter.
Sequence probability is the chain of next-token probabilities.
Next: Next Token Prediction—the training and generation objective itself.

Trainer’s Guide

Hands-on idea: Give students three hand-written logit vectors; they compute softmax by hand for a 4-word vocab, then verify with PyTorch.

Discussion prompt: When would you prefer sampling over always taking the mode? (Creative writing vs code completion with a clear syntax continuation.)

Recap: Next-token outputs are categorical distributions obtained via softmax over vocab logits. Continue with Next Token Prediction.