def step_fn(ids): # Token ids in, numpy logits for the next token out
logits = model(d2l.tensor([ids], ctx=d2l.try_gpu()))
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 I had come into the future. I was a vast array of a vast array of a vast array of a vast array of a vast array of a vast arr'
'it seemed to me that the sun willutes of a tramblance of the Time Traveller. The Time Traveller. The Time Traveller. The Time Traveller. The Time ...- 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.23, distinct-3 0.45: 'the time traveller I had come into the future. I was a vast array of a vast array of a vast array of a vast array of a v'
k=4, score -2.40, distinct-3 0.29: "the time traveller-side windows. ...
k=1, score -4.62, distinct-3 0.50: 'it seemed to me that the sun willutes of a tramblance of the Time Traveller. The Time Traveller. The Time Traveller. The ...
k=4, score -3.42, distinct-3 0.76: "it seemed to me that they were no doubt. I thought I heard a signs of the Time Traveller. The Time Traveller. ...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 0x76b93a316000>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.28, top-k keeps 20, top-p keeps 40, min-p keeps 12
'it seemed to me that': top prob 0.14, top-k keeps 20, top-p keeps 56, min-p keeps 20The 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.35): 'the time traveller I had come into the future. I was a vast array of a vast array of a vast array of a vast array of a vast array ...
T=1.0 (distinct-3 1.00): 'the time travelleriveley. (ason b another queuge and ta-haoces I scattered� lirial swific earlyt. Oittly comfort. I should a shome'
top-k (distinct-3 1.00): "the time traveller by moderonke at it for any more. 'I supposit and course. It is this lit a civiliz. I felt asked in avolution. ...
top-p (distinct-3 1.00): "the time traveller of a cigarened intoansites on ekingic were direain of the Time Traveller, however, catodous feet. 'Halloations ...
min-p (distinct-3 1.00): 'the time traveller-sidered that, was the chinased me. And there are of our own time. What it was this yet railway to the shaday ...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.