Where greedy decoding loses, and why top-p replaced top-k.
The notebook carries the source, the install step and every run below, so nothing has to be on your machine. Open it, then choose Runtime › Run all.
Colab opens it read only. Click Copy to Drive to keep
your changes. You can also download
decoding.ipynb, or run it
locally: python3 decoding.py.
python3 decoding.py
Every strategy on the same distribution.
python3 decoding.py --beam
Greedy takes the at 0.40 over a at 0.35 and never sees that a leads to 0.2992. It returns 0.1320. A beam of width 2 finds the better sequence.
python3 decoding.py --truncate
At k = 5 the same setting keeps 99 per cent of a peaked distribution and 64 per cent of a flat one. At p = 0.9 the candidate set follows the model's own confidence.
Download decoding.py
· served verbatim at https://nlp.jcrlabz.com/code/decoding.py
"""Decoding strategies, worked as Chapter 18 works them.
Four things, all reproducing the chapter's tables:
greedy against beam where greedy loses, and what beam costs
temperature one knob, three behaviours
top-k and top-p why a fixed k fails on some distributions
the whole picture the same distribution under every strategy
python3 decoding.py # all four
python3 decoding.py --beam # greedy against beam search
python3 decoding.py --temp # the temperature sweep
python3 decoding.py --truncate # top-k against top-p
python3 decoding.py --compare # every strategy, side by side
Install: nothing, the Python standard library is enough
"""
import argparse
import math
# A tiny language model given as an explicit tree of next-token probabilities.
# The key is the prefix; the value maps next token to probability.
TREE = {
(): {"the": 0.40, "a": 0.35, "one": 0.25},
("the",): {"cat": 0.55, "dog": 0.45},
("a",): {"cat": 0.10, "bird": 0.90},
("one",): {"cat": 0.50, "dog": 0.50},
("the", "cat"): {"sat": 0.60, "ran": 0.40},
("the", "dog"): {"sat": 0.50, "ran": 0.50},
("a", "cat"): {"sat": 0.50, "ran": 0.50},
("a", "bird"): {"sang": 0.95, "flew": 0.05},
("one", "cat"): {"sat": 0.50, "ran": 0.50},
("one", "dog"): {"sat": 0.50, "ran": 0.50},
}
# A separate distribution for the truncation demonstrations. Two shapes: one
# where the model is confident and one where it is not.
PEAKED = {"the": 0.85, "a": 0.07, "one": 0.04, "some": 0.02,
"any": 0.01, "each": 0.005, "every": 0.005}
FLAT = {"red": 0.14, "blue": 0.13, "green": 0.13, "black": 0.12,
"white": 0.12, "grey": 0.12, "brown": 0.12, "pink": 0.12}
def greedy(depth=3):
prefix, logp, trace = (), 0.0, []
for _ in range(depth):
dist = TREE.get(prefix)
if not dist:
break
w = max(dist, key=dist.get)
trace.append((prefix, w, dist[w], dict(dist)))
logp += math.log(dist[w])
prefix = prefix + (w,)
return prefix, math.exp(logp), trace
def all_sequences(depth=3):
out = []
def walk(prefix, p):
dist = TREE.get(prefix)
if not dist or len(prefix) == depth:
out.append((prefix, p))
return
for w, q in dist.items():
walk(prefix + (w,), p * q)
walk((), 1.0)
return sorted(out, key=lambda x: -x[1])
def beam(width, depth=3):
beams = [((), 0.0)]
for _ in range(depth):
cand = []
for prefix, lp in beams:
dist = TREE.get(prefix)
if not dist:
cand.append((prefix, lp))
continue
for w, q in dist.items():
cand.append((prefix + (w,), lp + math.log(q)))
cand.sort(key=lambda x: -x[1])
beams = cand[:width]
return [(p, math.exp(lp)) for p, lp in beams]
# --------------------------------------------------------------- beam search
def show_beam():
print("\nSTEP 1 greedy takes the best word and loses the best sentence\n")
seq, p, trace = greedy()
print(" GREEDY. At each step take the highest probability token.\n")
print(f" {'prefix':<18}{'choices':<34}{'taken':<8}{'p':>7}")
print(" " + "-" * 68)
for prefix, w, q, dist in trace:
pre = " ".join(prefix) if prefix else "(start)"
ch = ", ".join(f"{k} {v:.2f}" for k, v in dist.items())
print(f" {pre:<18}{ch:<34}{w:<8}{q:>7.2f}")
print(f"\n greedy output: '{' '.join(seq)}' probability {p:.4f}\n")
ranked = all_sequences()
print(" Now every three-word sequence, ranked by probability:\n")
print(f" {'rank':>5} {'sequence':<24}{'probability':>13}")
print(" " + "-" * 46)
for i, (s, q) in enumerate(ranked[:6], start=1):
mark = " <- greedy found this" if s == seq else ""
print(f" {i:>5} {' '.join(s):<24}{q:>13.4f}{mark}")
best = ranked[0]
print(f"\n The best sequence is '{' '.join(best[0])}' at {best[1]:.4f}.")
print(f" Greedy returned {p:.4f}, which is"
f" {(best[1] / p - 1) * 100:.0f} per cent worse.")
print(f"\n Greedy went wrong at the very first step. It took 'the' at")
print(f" 0.40 over 'a' at 0.35, and never saw that 'a' leads to 'bird'")
print(f" at 0.90 and then 'sang' at 0.95.")
print(f"\n One high-probability token can sit in front of a low-")
print(f" probability continuation. Greedy cannot know that in advance.")
print(f"\n\n BEAM SEARCH keeps the best w partial sequences at each step.\n")
print(f" {'width':>6} {'best sequence found':<26}{'probability':>13}"
f"{'optimal?':>10}")
print(" " + "-" * 58)
for w in (1, 2, 3, 5):
b = beam(w)
top = b[0]
ok = "yes" if abs(top[1] - best[1]) < 1e-12 else "no"
print(f" {w:>6} {' '.join(top[0]):<26}{top[1]:>13.4f}{ok:>10}")
print(f"\n Width 1 is greedy by definition. Width 2 already finds the")
print(f" optimum here, and wider beams cost more for nothing.")
print(f"\n Beam search is not guaranteed to find the best sequence. It")
print(f" searches more of the space than greedy and less than all of it,")
print(f" which is the only honest description of it.")
# --------------------------------------------------------------- temperature
def apply_temp(dist, T):
if T == 0:
top = max(dist, key=dist.get)
return {k: (1.0 if k == top else 0.0) for k in dist}
z = {k: math.exp(math.log(v) / T) for k, v in dist.items()}
s = sum(z.values())
return {k: v / s for k, v in z.items()}
def entropy(dist):
h = -sum(p * math.log2(p) for p in dist.values() if p > 0)
return abs(h) # avoid printing -0.0000 when one token has all the mass
def show_temp():
print("\n\nSTEP 2 temperature, one knob\n")
print(" Divide the logits by T before the softmax. Equivalently, raise")
print(" each probability to the power 1/T and renormalise.\n")
base = PEAKED
words = list(base)
print(f" {'T':>6}" + "".join(f"{w:>9}" for w in words[:5])
+ f"{'entropy':>10}")
print(" " + "-" * 62)
for T in (0.0, 0.5, 0.8, 1.0, 1.5, 2.0):
d = apply_temp(base, T)
cells = "".join(f"{d[w]:>9.4f}" for w in words[:5])
print(f" {T:>6.1f}{cells}{entropy(d):>10.4f}")
print(f"\n T = 1 leaves the model's own distribution alone.")
print(f" T below 1 sharpens it. At T = 0 it is greedy, and the entropy")
print(f" is 0 because there is nothing left to choose.")
print(f" T above 1 flattens it and raises the entropy.")
print(f"\n The trade has a name at each end. Low temperature gives")
print(f" repetitive, safe text. High temperature gives varied text that")
print(f" drifts off topic and eventually stops making sense.")
# --------------------------------------------------------- top-k and top-p
def top_k(dist, k):
keep = sorted(dist.items(), key=lambda x: -x[1])[:k]
s = sum(v for _, v in keep)
return {w: v / s for w, v in keep}
def top_p(dist, p):
keep, run = [], 0.0
for w, v in sorted(dist.items(), key=lambda x: -x[1]):
keep.append((w, v))
run += v
if run >= p:
break
s = sum(v for _, v in keep)
return {w: v / s for w, v in keep}
def show_truncate():
print("\n\nSTEP 3 top-k has a problem that top-p does not\n")
print(" Two next-token distributions. One where the model is sure, one")
print(" where it genuinely is not.\n")
for name, d in (("peaked", PEAKED), ("flat", FLAT)):
top = sorted(d.items(), key=lambda x: -x[1])[:4]
cells = ", ".join(f"{w} {v:.2f}" for w, v in top)
print(f" {name:<8}{cells}, ... entropy {entropy(d):.4f} bits")
print(f"\n TOP-K keeps a fixed number of candidates, whatever the shape.\n")
print(f" {'k':>4}{'peaked: kept':>16}{'mass kept':>12}"
f"{'flat: kept':>14}{'mass kept':>12}")
print(" " + "-" * 60)
for k in (1, 2, 5, 8):
pk = sorted(PEAKED.items(), key=lambda x: -x[1])[:k]
fk = sorted(FLAT.items(), key=lambda x: -x[1])[:k]
print(f" {k:>4}{len(pk):>16}{sum(v for _, v in pk):>12.4f}"
f"{len(fk):>14}{sum(v for _, v in fk):>12.4f}")
print(f"\n Read k = 5. On the peaked distribution it admits four tokens")
print(f" the model had all but ruled out, together worth"
f" {sum(v for _,v in sorted(PEAKED.items(), key=lambda x:-x[1])[1:5]):.3f}.")
print(f" On the flat one it discards three tokens that were as good as")
print(f" the ones it kept.")
print(f"\n The same k is too loose on one shape and too tight on the")
print(f" other, because k cannot see the shape.")
print(f"\n\n TOP-P keeps the smallest set whose mass reaches p.\n")
print(f" {'p':>6}{'peaked: kept':>16}{'mass':>10}{'flat: kept':>14}"
f"{'mass':>10}")
print(" " + "-" * 58)
for p in (0.5, 0.9, 0.95, 0.99):
pk, fk = top_p(PEAKED, p), top_p(FLAT, p)
mp = sum(PEAKED[w] for w in pk)
mf = sum(FLAT[w] for w in fk)
print(f" {p:>6.2f}{len(pk):>16}{mp:>10.4f}{len(fk):>14}{mf:>10.4f}")
print(f"\n At p = 0.9 the peaked distribution keeps"
f" {len(top_p(PEAKED, 0.9))} tokens and the flat one keeps"
f" {len(top_p(FLAT, 0.9))}.")
print(f" The size of the candidate set now follows the model's own")
print(f" confidence, which is exactly what a fixed k could not do.")
print(f"\n That is why nucleus sampling replaced top-k as the default.")
# ------------------------------------------------------------ side by side
def show_compare():
print("\n\nSTEP 4 the same distribution under every strategy\n")
d = PEAKED
rows = [
("greedy", apply_temp(d, 0.0)),
("temperature 0.7", apply_temp(d, 0.7)),
("pure sampling, T = 1", dict(d)),
("temperature 1.5", apply_temp(d, 1.5)),
("top-k, k = 3", top_k(d, 3)),
("top-p, p = 0.9", top_p(d, 0.9)),
]
words = list(d)
print(f" {'strategy':<22}" + "".join(f"{w[:6]:>8}" for w in words)
+ f"{'kept':>7}{'entropy':>9}")
print(" " + "-" * 86)
for name, dd in rows:
cells = "".join(f"{dd.get(w, 0.0):>8.3f}" for w in words)
kept = sum(1 for w in words if dd.get(w, 0.0) > 0)
print(f" {name:<22}{cells}{kept:>7}{entropy(dd):>9.4f}")
print(f"\n Every row is the same model. None of them changed a weight.")
print(f" Decoding is a decision made after the model has spoken, and it")
print(f" changes the output more than most fine-tuning does.")
print(f"\n Which to use follows from the task. A factual answer wants")
print(f" the low-entropy rows. A story wants one of the middle ones.")
print(f" Nothing wants the bottom of the temperature range or the top.")
def main():
ap = argparse.ArgumentParser(description=__doc__)
for f in ("beam", "temp", "truncate", "compare"):
ap.add_argument(f"--{f}", action="store_true")
a = ap.parse_args()
picked = a.beam or a.temp or a.truncate or a.compare
if a.beam or not picked:
show_beam()
if a.temp or not picked:
show_temp()
if a.truncate or not picked:
show_truncate()
if a.compare or not picked:
show_compare()
print()
if __name__ == "__main__":
main()