def step_fn(ids): # Token ids in, numpy logits for the next token out
logits = model(d2l.tensor([ids]))
return d2l.numpy(logits)[0, -1]Decoding and Generation
Decoding and Generation
A trained LM supplies conditionals P(x_t \mid x_{<t}). Decoding turns them into actual sequences, and the strategy is a design decision of its own:
- Maximization (greedy, beam search): tasks with one right answer.
- Sampling (temperature, top-k/p, min-p): open-ended generation.
The decoding problem
The chain rule scores any full sequence:
P(x_1, \ldots, x_T) = \prod_{t=1}^{T} P(x_t \mid x_{<t}).
The “best” sequence is an argmax over |\mathcal{V}|^T candidates: 1024^{50} \approx 10^{150}. Exhaustive search is hopeless.
Worse: for open-ended text the argmax is not even desirable. Many acceptable continuations, each individually improbable; picking the single most probable one discards the diversity we asked for.
One narrow interface
Train the RNN LM of 9.5 (same recipe), then wrap it: token ids in, numpy logits for the next token out. Every strategy below is plain Python against this interface:
Works for any model that predicts next tokens: this RNN, an LSTM, a transformer.
Greedy decoding
x_t = \operatorname*{argmax}_{x \in \mathcal{V}} P(x \mid x_{<t})
'the time traveller the\nside, I saw\nlittleess of the\nfour of the floor of the\nsunned and\npresented\ntheseength, the\nfour of the faint of'
'it seemed to me that the\nhim, and that the\nfour of the floor of the\nsunned and\npresented\ntheseength, the\nfour of the faint of\nthese, I'- Myopic: locally best tokens \neq best sequence.
- Mode-seeking: argmax feedback amplifies stock phrases; repetition is self-reinforcing (neural text degeneration).
Beam search
Keep the k best prefixes alive; expand each over the vocabulary, keep the best k of the k \cdot |\mathcal{V}| candidates:
Cost \mathcal{O}(k |\mathcal{V}| T): greedy is k = 1, exhaustive is k \to \infty.
Length normalization
Log-probabilities only fall as sequences grow, so raw sums favor short candidates. Rank by
\frac{1}{L^\alpha} \sum_{t=m+1}^{m+L} \log P(x_t \mid x_{<t}), \qquad \alpha \approx 0.75.
def beam_search(step_fn, prefix, num_tokens, beam_size=4,
alpha=0.75, eos_id=None):
"""Length-normalized beam search (score = log P / len^alpha).
Returns candidates sorted best-first; [0][1] is the argmax sequence."""
def log_softmax(logits):
logits = [float(l) for l in logits]
m = max(logits)
lse = m + math.log(sum(math.exp(l - m) for l in logits))
return [l - lse for l in logits]
score = lambda logp, ids: logp / max(1, len(ids) - len(prefix))**alpha
beams = [(0.0, list(prefix), False)] # (log-prob, sequence, finished)
for _ in range(num_tokens):
if all(done for _, _, done in beams):
break
candidates = [b for b in beams if b[2]] # Finished pass through
for logp, ids, _ in (b for b in beams if not b[2]):
logprobs = log_softmax(step_fn(ids))
best = sorted(range(len(logprobs)),
key=lambda i: -logprobs[i])[:beam_size]
candidates += [(logp + logprobs[i], ids + [i], i == eos_id)
for i in best]
beams = sorted(candidates, key=lambda b: score(b[0], b[1]),
reverse=True)[:beam_size]
return sorted([(score(logp, ids), ids) for logp, ids, _ in beams],
reverse=True)Better score, worse text
def distinct(ids, n=3):
grams = [tuple(ids[i:i + n]) for i in range(len(ids) - n + 1)]
return len(set(grams)) / max(1, len(grams))
for text in ('the time traveller', 'it seemed to me that'):
pre = data.tokenizer.encode(text)
for k in (1, 4):
score, ids = beam_search(step_fn, pre, 40, beam_size=k)[0]
print(f'k={k}, score {score:.2f}, distinct-3 '
f'{distinct(ids[len(pre):]):.2f}: '
f'{data.tokenizer.decode(ids)!r}')k=1, score -4.74, distinct-3 0.95: 'the time traveller the\nside, I saw\nlittleess of the\nfour of the floor of the\nsunned and\npresented\ntheseength'
k=4, score -3.16, distinct-3 0.45: "the time traveller the Time Traveller.\n\n'There I found a\npresented.\n\n'There I found a\npresented.\n\n'There I found a"
k=1, score -4.62, distinct-3 0.95: 'it seemed to me that the\nhim, and that the\nfour of the floor of the\nsunned and\npresented\ntheseength, the\nfour'
k=4, score -3.15, distinct-3 0.58: "it seemed to me that I had\ndisappef.\n\n'There I found a\npresented.\n\n'There I found a\npresented.\n\n'"k = 4 usually finds sequences the model scores higher (it is an approximate search, so not always), and those are the ones that lock into a literal loop (distinct-3 collapses). For open-ended text the argmax is the wrong target; search cannot fix that. Large beams even hurt MT. Beam search survives where its assumptions hold: ASR, MT, constrained decoding.
Sampling and the unreliable tail
Draw x_t \sim P(x_t \mid x_{<t}): diverse, no repetition traps. But the tail’s combined mass is large, and one absurd token derails the context.
Temperature rescales logits:
P_T(x \mid x_{<t}) \propto P(x \mid x_{<t})^{1/T}
T \to 0 greedy, T = 1 the model, T \to \infty uniform; head and tail reshape together.
Truncating the tail
- top-k: keep the k most probable (fixed count).
- top-p (nucleus): smallest head with cumulative mass \geq p (fixed mass).
- min-p: keep P(x) \geq p_{\min} \cdot \max_x P(x) (relative to the model’s own confidence; robust at high temperature).
A unified sampler
def sample_next(logits, strategy='greedy', temperature=1.0,
k=None, p=None, min_p=None, rng=None):
"""Choose the next token id from a 1-D numpy logits array.
strategy: 'greedy' | 'sample' (with optional top-k / top-p / min-p
truncation applied to the temperature-scaled distribution)."""
logits = [float(l) for l in logits]
if strategy == 'greedy' or temperature == 0:
return max(range(len(logits)), key=lambda i: logits[i])
m = max(logits)
probs = [math.exp((l - m) / temperature) for l in logits]
total = sum(probs)
probs = [q / total for q in probs]
order = sorted(range(len(probs)), key=lambda i: -probs[i])
keep = len(order)
if k is not None: # Top-k: the k most probable tokens
keep = min(keep, k)
if p is not None: # Top-p: smallest head with cumulative mass >= p
mass, n = 0.0, 0
while mass < p and n < len(order):
mass, n = mass + probs[order[n]], n + 1
keep = min(keep, n)
if min_p is not None: # Min-p: within a factor of the top token
bar = min_p * probs[order[0]]
keep = min(keep, sum(q >= bar for q in probs))
kept = order[:keep]
rng = random if rng is None else rng
r = rng.random() * sum(probs[i] for i in kept)
for i in kept:
r -= probs[i]
if r <= 0:
return i
return kept[-1]Generation is a loop
def generate(step_fn, prefix, num_tokens, eos_id=None, **strategy):
"""Autoregressive generation. step_fn(ids: list[int]) -> numpy logits
for the next token. Returns prefix + continuation (stops on eos_id)."""
ids = list(prefix)
for _ in range(num_tokens):
ids.append(sample_next(step_fn(ids), **strategy))
if eos_id is not None and ids[-1] == eos_id:
break
return idsTruncation on a toy distribution (0.45, 0.25, 0.15, 0.10, 0.05):
greedy: 0
{'k': 2} [0, 0, 0, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1, 0, 1]
{'p': 0.85} [0, 2, 1, 0, 0, 0, 0, 1, 1, 1, 0, 2, 2, 1, 1]
{'min_p': 0.5} [1, 0, 0, 1, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0]One distribution, three cutoffs
def next_probs(text): # Sorted next-token distribution after a prefix
logits = step_fn(data.tokenizer.encode(text))
probs = np.exp(logits - logits.max())
return np.sort(probs / probs.sum())[::-1]
probs = next_probs('the time traveller')
n_p = int((np.cumsum(probs) < 0.9).sum()) + 1
n_m = int((probs >= 0.05 * probs[0]).sum())
d2l.set_figsize((5.5, 3))
d2l.plt.loglog(np.arange(1, len(probs) + 1), probs, color='C0')
d2l.plt.axvline(20, color='C1', ls='--', label='top-k, k=20')
d2l.plt.axvline(n_p, color='C2', ls='-.', label=f'top-p=0.9: keep {n_p}')
d2l.plt.axhline(0.05 * probs[0], color='C3', ls=':',
label=f'min-p=0.05: keep {n_m}')
d2l.plt.xlabel('token rank')
d2l.plt.ylabel('probability')
d2l.plt.legend()
<matplotlib.legend.Legend at 0x73b5d87d9880>top-k: vertical cut at fixed rank; top-p: vertical cut that slides with the head; min-p: horizontal cut that rides the model’s confidence.
The dials adapt (or don’t)
'the time traveller': top prob 0.07, top-k keeps 20, top-p keeps 149, min-p keeps 42
'it seemed to me that': top prob 0.20, top-k keeps 20, top-p keeps 53, min-p keeps 11The more confident the model, the smaller the kept sets, and min-p tightens the fastest; top-k keeps 20 no matter what.
The same prefix under every strategy
strategies = {'greedy': dict(strategy='greedy'),
'T=1.0': dict(strategy='sample'),
'top-k': dict(strategy='sample', k=20),
'top-p': dict(strategy='sample', p=0.9),
'min-p': dict(strategy='sample', min_p=0.05)}
for name, s in strategies.items():
out = generate(step_fn, prefix, 50, rng=np.random.default_rng(0), **s)
print(f'{name:>7} (distinct-3 {distinct(out[len(prefix):]):.2f}): '
f'{data.tokenizer.decode(out)!r}')greedy (distinct-3 0.85): 'the time traveller the\nside, I saw\nlittleess of the\nfour of the floor of the\nsunned and\npresented\ntheseength, the\nfour of ...
T=1.0 (distinct-3 1.00): 'the time travellers and.\nI either him to and approumpessity upon me. It was\nbeto either� lever of proenox orually the twoworic, ...
top-k (distinct-3 1.00): "the time travellerces, and I went downsetulate\nMy and\ns the breathing of detewomion, it's have been during. Soets and spiming. ...
top-p (distinct-3 1.00): "the time travellercondon of white had come me thinks upon\nthated; and eitherance of\nbacstand. So who Man oner, and ...
min-p (distinct-3 1.00): "the time traveller\nthereenumy of a marms of us. 'I bewered.\n\n'Tondeedorce the\nwellorner, the night, and that was the starteroots!"Greedy: generic and stuck. Pure sampling: diverse, erratic. Truncated sampling: the useful middle, and the default in deployed systems.
Recap
- Decoding = choosing sequences from exponentially many using local conditionals; maximization vs sampling.
- Greedy is myopic; beam search fixes myopia at k\times cost, with length normalization, and increases repetition on open-ended text.
- Sampling dials: temperature, top-k, top-p, min-p.
- Perplexity scores the model; judging generated text needs humans or proxies.
- Naive generation re-runs the model per token: \mathcal{O}(T^2). A recurrent state makes it \mathcal{O}(1) per token; next chapter builds on exactly this.