How a language model works all 5 written
Tokens, next-word prediction, attention and a transformer block, from nothing.
17 Tokenization, and why it explains so many bugs
A model never sees your text. It sees a list of integers, and the thing that produced them is a lookup table built by counting character pairs in a corpus that was mostly English. That table is frozen before training starts. Every quirk in this lesson — letter counting, reversing strings, digits, the price of writing in Greek — is a consequence of it, not of the neural network at all.
What byte pair encoding actually does
Start with every word spelled out as single characters. Count every adjacent pair. Merge the most frequent pair into one new symbol. Repeat. That is the whole algorithm, and the merges it learns are the vocabulary.
# Byte pair encoding, trained from nothing on a tiny corpus.
CORPUS = ("low low low low low lowest lowest "
"newer newer newer newer newer newer "
"wider wider wider new new")
def word_freqs(text):
freqs = {}
for word in text.split():
freqs[word] = freqs.get(word, 0) + 1
return freqs
def start_splits(freqs):
"""Every word starts life as single characters, plus an end-of-word marker."""
return {tuple(word) + ("_",): n for word, n in freqs.items()}
def pair_counts(splits):
counts = {}
for symbols, n in splits.items():
for a, b in zip(symbols, symbols[1:]):
counts[(a, b)] = counts.get((a, b), 0) + n
return counts
def apply_merge(symbols, a, b):
"""Replace every adjacent a, b with the single symbol a + b."""
out = []
i = 0
while i < len(symbols):
if i + 1 < len(symbols) and symbols[i] == a and symbols[i + 1] == b:
out.append(a + b)
i += 2
else:
out.append(symbols[i])
i += 1
return tuple(out)
def train_bpe(text, rounds):
splits = start_splits(word_freqs(text))
merges = []
for _ in range(rounds):
counts = pair_counts(splits)
if not counts:
break
# sorted() first, so a tie is broken alphabetically rather than by dict order.
# Without it, the vocabulary depends on insertion order and is not reproducible.
best = max(sorted(counts), key=lambda pair: counts[pair])
merges.append((best, counts[best]))
splits = {apply_merge(s, best[0], best[1]): n for s, n in splits.items()}
return merges, splits
merges, splits = train_bpe(CORPUS, 10)
print("corpus words:", word_freqs(CORPUS))
print()
print(" round merged pair seen new symbol")
for i, ((a, b), n) in enumerate(merges, start=1):
print(" {:5} {:6} + {:6} {:4} {}".format(i, repr(a), repr(b), n, repr(a + b)))
print()
print("how each word is spelled after ten merges:")
for symbols, n in sorted(splits.items()):
print(" {:3} x {}".format(n, list(symbols)))
frequencies = [n for _, n in merges]
print()
print("vocabulary:", sorted(set(sym for symbols in splits for sym in symbols)))
test("the first merge is the most frequent adjacent pair in the corpus", merges[0][0], ("e", "r"))
test("ten rounds learn exactly ten merges", len(merges), 10)
test("each merge is at least as frequent as the one after it",
frequencies, sorted(frequencies, reverse=True))
test("the most common word in the corpus ends up as one symbol", ("newer_",) in splits, True)
test("a rarer word is still several", ("low", "e", "s", "t", "_") in splits, True)What it printed, when the build ran it
corpus words: {'low': 5, 'lowest': 2, 'newer': 6, 'wider': 3, 'new': 2}
round merged pair seen new symbol
1 'e' + 'r' 9 'er'
2 'er' + '_' 9 'er_'
3 'e' + 'w' 8 'ew'
4 'n' + 'ew' 8 'new'
5 'l' + 'o' 7 'lo'
6 'lo' + 'w' 7 'low'
7 'new' + 'er_' 6 'newer_'
8 'low' + '_' 5 'low_'
9 'd' + 'er_' 3 'der_'
10 'i' + 'der_' 3 'ider_'
how each word is spelled after ten merges:
2 x ['low', 'e', 's', 't', '_']
5 x ['low_']
2 x ['new', '_']
6 x ['newer_']
3 x ['w', 'ider_']
vocabulary: ['_', 'e', 'ider_', 'low', 'low_', 'new', 'newer_', 's', 't', 'w']
ok the first merge is the most frequent adjacent pair in the corpus
ok ten rounds learn exactly ten merges
ok each merge is at least as frequent as the one after it
ok the most common word in the corpus ends up as one symbol
ok a rarer word is still several
Read the merge table top to bottom and you are reading the vocabulary being built.
e plus r first, because newer and wider
between them use it nine times. Then er plus the end marker, which is what makes
the suffix -er a unit. By round seven the whole word newer is a single
symbol, because the corpus contains it six times. Frequency buys you a short encoding, and
nothing else does.
The end-of-word marker matters more than it looks. Without it, the symbol
low learned from low and the low inside
lowest are the same thing, and the model cannot tell a whole word from a
prefix. Real tokenizers do the same job by attaching the leading space to the
token, which is why " cat" and "cat" are two different tokens
with two different ids — and why a prompt that ends in a trailing space quietly changes
what the model predicts next.
Encoding: replay the merges, in order
Training produced an ordered list. Encoding a new word means starting from characters and replaying that list, in exactly the order it was learned. The order is part of the tokenizer.
# Encoding with the merges we just learned. Same list, same order, every time.
CORPUS = ("low low low low low lowest lowest "
"newer newer newer newer newer newer "
"wider wider wider new new")
MERGES = [("e", "r"), ("er", "_"), ("e", "w"), ("n", "ew"), ("l", "o"),
("lo", "w"), ("new", "er_"), ("low", "_"), ("d", "er_"), ("i", "der_")]
def apply_merge(symbols, a, b):
out = []
i = 0
while i < len(symbols):
if i + 1 < len(symbols) and symbols[i] == a and symbols[i + 1] == b:
out.append(a + b)
i += 2
else:
out.append(symbols[i])
i += 1
return tuple(out)
def encode(word, merges):
"""Characters first, then replay every merge in the order it was learned."""
symbols = tuple(word) + ("_",)
for a, b in merges:
symbols = apply_merge(symbols, a, b)
return list(symbols)
for word in ("low", "newer", "lower", "newest", "widen", "glow"):
pieces = encode(word, MERGES)
seen = "in the corpus" if word in CORPUS.split() else "never seen"
print(" {:8} -> {:38} {} token(s) {}".format(word, str(pieces), len(pieces), seen))
print()
print("Order is part of the tokenizer. Replay the merges backwards and you get")
print("a different spelling of the same word:")
print(" newer, merges in order :", encode("newer", MERGES))
print(" newer, merges reversed :", encode("newer", list(reversed(MERGES))))
frequent = encode("newer", MERGES)
rare = encode("newest", MERGES)
test("a word the tokenizer saw six times costs one token", len(frequent), 1)
test("a word it never saw costs more", len(rare) > len(frequent), True)
test("lower is not a token, it is low plus the -er suffix",
encode("lower", MERGES), ["low", "er_"])
test("nothing is lost, the pieces still spell the word",
"".join(encode("newest", MERGES)).replace("_", ""), "newest")
test("replaying the merges out of order changes the answer",
encode("newer", MERGES) == encode("newer", list(reversed(MERGES))), False)
test("a character the tokenizer never saw survives as itself",
"g" in encode("glow", MERGES), True)What it printed, when the build ran it
low -> ['low_'] 1 token(s) in the corpus newer -> ['newer_'] 1 token(s) in the corpus lower -> ['low', 'er_'] 2 token(s) never seen newest -> ['new', 'e', 's', 't', '_'] 5 token(s) never seen widen -> ['w', 'i', 'd', 'e', 'n', '_'] 6 token(s) never seen glow -> ['g', 'low_'] 2 token(s) never seen Order is part of the tokenizer. Replay the merges backwards and you get a different spelling of the same word: newer, merges in order : ['newer_'] newer, merges reversed : ['n', 'ew', 'er', '_'] ok a word the tokenizer saw six times costs one token ok a word it never saw costs more ok lower is not a token, it is low plus the -er suffix ok nothing is lost, the pieces still spell the word ok replaying the merges out of order changes the answer ok a character the tokenizer never saw survives as itself
That table is the practical shape of every tokenizer you will meet. Common words are one token. Anything unusual — a surname, a product code, a typo, an identifier from your codebase — falls apart into pieces, and you pay per piece. It is also why the same paragraph rewritten in plainer words is genuinely cheaper to send.
Why it cannot count the letters in "strawberry"
Here is the payoff. The model receives integers. The letters inside each integer are not
available to it any more than the pixels of this page are available to you as voltages. To
count the r in a word, the model has to have memorised the spelling of each
token id, in its weights, from text.
# What the model actually receives: integers. Letters are not in there.
# This vocabulary is hand-written to match the way a real tokenizer splits
# these words, and applied greedily, longest piece first.
VOCAB = {
"str": 2, "aw": 3, "berry": 4, "berries": 5, "count": 6,
"the": 7, "letters": 8, "in": 9, " ": 10,
"a": 20, "b": 21, "e": 22, "h": 23, "i": 24, "r": 25, "s": 26, "t": 27,
"w": 28, "y": 29, "d": 30, "o": 31, "n": 32, "c": 33, "l": 34, "u": 35,
}
def encode(text, vocab):
"""Greedy longest match. The real thing is fancier; the effect is the same."""
pieces = []
i = 0
while i < len(text):
for size in range(min(8, len(text) - i), 0, -1):
piece = text[i:i + size]
if piece in vocab:
pieces.append(piece)
i += size
break
else:
raise ValueError("no token covers {!r}".format(text[i]))
return pieces
WORD = "strawberry"
pieces = encode(WORD, VOCAB)
ids = [VOCAB[p] for p in pieces]
print("the word :", WORD, " ", len(WORD), "letters")
print("its pieces :", pieces)
print("what the model gets :", ids)
print()
print("letter r in the word :", WORD.count("r"))
print("pieces containing an r :", sum(1 for p in pieces if "r" in p))
print("A count done over the ids answers 2. The answer is 3.")
print()
print("reversing, two ways:")
print(" reverse the letters :", WORD[::-1])
print(" reverse the tokens :", "".join(reversed(pieces)))
print("A model rearranges tokens. Only one of those lines is a reversal.")
letters = WORD.count("r")
per_piece = sum(1 for p in pieces if "r" in p)
test("ten letters arrive as three tokens", [len(WORD), len(pieces)], [10, 3])
test("not one of those tokens is a single letter", all(len(p) > 1 for p in pieces), True)
test("the word really does contain three r", letters, 3)
test("counting over tokens instead of letters gives the wrong answer",
per_piece == letters, False)
test("reversing the token order does not reverse the word",
"".join(reversed(pieces)) == WORD[::-1], False)
test("the reversal a reader wants needs the spelling, which the ids do not carry",
WORD[::-1], "yrrebwarts")
test("the pieces do still spell the word, if you look each id up", "".join(pieces), WORD)What it printed, when the build ran it
the word : strawberry 10 letters its pieces : ['str', 'aw', 'berry'] what the model gets : [2, 3, 4] letter r in the word : 3 pieces containing an r : 2 A count done over the ids answers 2. The answer is 3. reversing, two ways: reverse the letters : yrrebwarts reverse the tokens : berryawstr A model rearranges tokens. Only one of those lines is a reversal. ok ten letters arrive as three tokens ok not one of those tokens is a single letter ok the word really does contain three r ok counting over tokens instead of letters gives the wrong answer ok reversing the token order does not reverse the word ok the reversal a reader wants needs the spelling, which the ids do not carry ok the pieces do still spell the word, if you look each id up
So the letter-counting failure is not a reasoning failure. The model is being asked a
question about a representation it does not have. It answers from a statistical memory of how
often text near the token berry mentioned the letter r, which is a
terrible way to count and works about as well as you would expect. The same argument explains
reversing a string, spelling a word backwards, counting words in a paragraph, and finding the
third character of a name.
This is also why the fix works. Ask the model to write the word out with spaces between
the letters first — s t r a w b e r r y — and each letter becomes its own
token, because that spacing does not appear in the merge table. Now the count is a
count over tokens, which is something the model can do. "Think step by step" earns its
reputation partly by accident: it converts problems about characters into problems about
tokens.
Numbers, and why arithmetic is harder than it should be
The tokenizer has no idea what a number is. It merged whatever digit strings were frequent in its corpus, so numbers of the same length can arrive as wildly different token counts, and the digits of two numbers you want to add do not line up.
# Numbers. The tokenizer never learned arithmetic, only what it saw often.
# A corpus where some numbers are common and others are not: the shape of the
# web text a real tokenizer is fitted to.
CORPUS = ("2020 2020 2020 2020 2020 2020 2020 2020 "
"2019 2019 2019 2019 2019 "
"100 100 100 100 100 100 "
"10 10 10 10 12 12 12 ")
def word_freqs(text):
freqs = {}
for word in text.split():
freqs[word] = freqs.get(word, 0) + 1
return freqs
def pair_counts(splits):
counts = {}
for symbols, n in splits.items():
for a, b in zip(symbols, symbols[1:]):
counts[(a, b)] = counts.get((a, b), 0) + n
return counts
def apply_merge(symbols, a, b):
out = []
i = 0
while i < len(symbols):
if i + 1 < len(symbols) and symbols[i] == a and symbols[i + 1] == b:
out.append(a + b)
i += 2
else:
out.append(symbols[i])
i += 1
return tuple(out)
def train_bpe(text, rounds):
splits = {tuple(w): n for w, n in word_freqs(text).items()}
merges = []
for _ in range(rounds):
counts = pair_counts(splits)
if not counts:
break
best = max(sorted(counts), key=lambda pair: counts[pair])
merges.append(best)
splits = {apply_merge(s, best[0], best[1]): n for s, n in splits.items()}
return merges
def encode(text, merges):
symbols = tuple(text)
for a, b in merges:
symbols = apply_merge(symbols, a, b)
return list(symbols)
MERGES = train_bpe(CORPUS, 6)
print("merges learned, in order:", ["".join(m) for m in MERGES])
print()
print(" number tokens how many")
for number in ("2020", "2019", "2021", "1000", "100", "137"):
pieces = encode(number, MERGES)
print(" {:8} {:30} {}".format(number, str(pieces), len(pieces)))
print()
print("Four digits cost one token or three, decided entirely by how often the")
print("tokenizer happened to see them. Column arithmetic needs digits lined up.")
print("These are not lined up, and the model has to learn the alignment itself.")
y2020 = encode("2020", MERGES)
y2021 = encode("2021", MERGES)
test("a year the tokenizer saw eight times is a single token", len(y2020), 1)
test("the next year along is not", len(y2021) > 1, True)
test("two numbers of the same length get different token counts",
len(y2020) == len(y2021), False)
test("so their digits do not line up with each other", y2020 == y2021, False)
test("nothing is corrupted, only chopped oddly", "".join(y2021), "2021")
test("a number it never saw is spelled out digit by digit",
encode("137", MERGES), ["1", "3", "7"])What it printed, when the build ran it
merges learned, in order: ['20', '10', '2020', '100', '19', '2019'] number tokens how many 2020 ['2020'] 1 2019 ['2019'] 1 2021 ['20', '2', '1'] 3 1000 ['100', '0'] 2 100 ['100'] 1 137 ['1', '3', '7'] 3 Four digits cost one token or three, decided entirely by how often the tokenizer happened to see them. Column arithmetic needs digits lined up. These are not lined up, and the model has to learn the alignment itself. ok a year the tokenizer saw eight times is a single token ok the next year along is not ok two numbers of the same length get different token counts ok so their digits do not line up with each other ok nothing is corrupted, only chopped oddly ok a number it never saw is spelled out digit by digit
This is why a model can be fluent about calculus and wrong about
4123 + 8991. It is not doing arithmetic on numbers, it is predicting tokens
after a sequence of unevenly chopped digit groups. It is also why newer tokenizers
deliberately split digits into fixed groups of one or three, and why the right engineering
answer is to hand the arithmetic to a tool rather than to the model. If a number matters —
a total, an id, a version — validate it, do not trust it.
Bytes, and why other languages cost more
A real tokenizer starts from UTF-8 bytes rather than characters, so it can encode anything without ever meeting an unknown symbol. But the merges were learned from a corpus that was overwhelmingly English. Text in a script that corpus barely contained falls back to almost one token per byte, and non-Latin scripts take two or three bytes per character.
# A real tokenizer works on UTF-8 bytes, and learned its merges from a corpus
# that was mostly English. Both halves of that sentence have a price.
ENGLISH_CORPUS = ("the cat sat on the mat and the cat sat on the hat " * 8 +
"the man on the mat sat and the cat ran on the sand " * 8)
def chunks(text):
"""Keep merges inside words. A real tokenizer does this with a regex, and
attaches the leading space to the word, exactly as this does."""
parts = text.split(" ")
return [parts[0]] + [" " + p for p in parts[1:]]
def to_symbols(text):
"""Start from raw UTF-8 bytes, so any text on earth can be encoded."""
return tuple((b,) for b in text.encode("utf-8"))
def pair_counts(sequences):
counts = {}
for symbols in sequences:
for a, b in zip(symbols, symbols[1:]):
counts[(a, b)] = counts.get((a, b), 0) + 1
return counts
def apply_merge(symbols, a, b):
out = []
i = 0
while i < len(symbols):
if i + 1 < len(symbols) and symbols[i] == a and symbols[i + 1] == b:
out.append(a + b)
i += 2
else:
out.append(symbols[i])
i += 1
return tuple(out)
def train(text, rounds):
sequences = [to_symbols(c) for c in chunks(text)]
merges = []
for _ in range(rounds):
counts = pair_counts(sequences)
if not counts:
break
best = max(sorted(counts), key=lambda pair: counts[pair])
merges.append(best)
sequences = [apply_merge(s, best[0], best[1]) for s in sequences]
return merges
def encode(text, merges):
out = []
for chunk in chunks(text):
symbols = to_symbols(chunk)
for a, b in merges:
symbols = apply_merge(symbols, a, b)
out.extend(symbols)
return out
def show(symbols):
"""Readable where the token is whole text, raw bytes where it is not."""
out = []
for s in symbols:
raw = bytes(s)
try:
out.append(raw.decode("utf-8"))
except UnicodeDecodeError:
out.append(repr(raw))
return out
MERGES = train(ENGLISH_CORPUS, 40)
SAMPLES = (
("English", "the cat sat on the mat"),
("Greek", "η γάτα κάθισε στο χαλί"),
("Japanese", "猫がマットの上に座った"),
)
print("40 merges, learned from English text only.")
print()
print(" language characters utf-8 bytes tokens tokens per character")
ratio = {}
for name, text in SAMPLES:
tokens = encode(text, MERGES)
ratio[name] = len(tokens) / len(text)
print(" {:9} {:10} {:11} {:6} {:.2f}".format(
name, len(text), len(text.encode("utf-8")), len(tokens), ratio[name]))
print()
print("English :", show(encode(SAMPLES[0][1], MERGES)))
print("Greek :", show(encode(SAMPLES[1][1], MERGES))[:8], "... and 32 more")
print("Japanese :", show(encode(SAMPLES[2][1], MERGES))[:6], "... and 27 more")
print()
print("Every Greek and Japanese token there is a single raw byte, because no")
print("merge in the table ever saw those bytes. Same sentence, three times the bill.")
rare = "the cat sat on the trampoline"
rare_tokens = encode(rare, MERGES)
print()
print("and token count is not word count:")
print(" ", rare)
print(" ", show(rare_tokens))
print(" {} words, {} tokens".format(len(rare.split()), len(rare_tokens)))
test("a common English word is one whole token", show(encode("the", MERGES)), ["the"])
test("the same sentence costs more per character in Greek than in English",
ratio["Greek"] > ratio["English"], True)
test("Japanese costs three tokens per character, one per UTF-8 byte",
close(ratio["Japanese"], 3.0), True)
test("token count is not word count", len(rare_tokens) == len(rare.split()), False)
test("one rare word is most of the bill for that sentence", len(rare_tokens) > 10, True)
test("nothing is lost, the bytes still decode to the original text",
bytes(b for s in encode(SAMPLES[1][1], MERGES) for b in s).decode("utf-8"),
SAMPLES[1][1])What it printed, when the build ran it
40 merges, learned from English text only. language characters utf-8 bytes tokens tokens per character English 22 22 6 0.27 Greek 22 40 40 1.82 Japanese 11 33 33 3.00 English : ['the', ' cat', ' sat', ' on', ' the', ' mat'] Greek : ["b'\\xce'", "b'\\xb7'", ' ', "b'\\xce'", "b'\\xb3'", "b'\\xce'", "b'\\xac'", "b'\\xcf'"] ... and 32 more Japanese : ["b'\\xe7'", "b'\\x8c'", "b'\\xab'", "b'\\xe3'", "b'\\x81'", "b'\\x8c'"] ... and 27 more Every Greek and Japanese token there is a single raw byte, because no merge in the table ever saw those bytes. Same sentence, three times the bill. and token count is not word count: the cat sat on the trampoline ['the', ' cat', ' sat', ' on', ' the', ' ', 't', 'r', 'a', 'm', 'p', 'o', 'l', 'i', 'n', 'e'] 6 words, 16 tokens ok a common English word is one whole token ok the same sentence costs more per character in Greek than in English ok Japanese costs three tokens per character, one per UTF-8 byte ok token count is not word count ok one rare word is most of the bill for that sentence ok nothing is lost, the bytes still decode to the original text
Read the ratio column and you have the answer to several practical questions at once. A context window measured in tokens holds roughly three times less Japanese than English. A per-token price is a higher per-sentence price for a Greek user. A chunk size of 500 tokens is a much smaller piece of text in Hindi than in French. None of that is a policy decision anyone made; it falls out of which corpus the merges were counted on.
- Cost and context. Budget in tokens, not words, and measure the ratio on your own real traffic rather than assuming the usual "one token is about four characters".
- Chunking for retrieval. A fixed token size cuts non-English text into much shorter passages. If retrieval quality drops for one language, check this first.
- Character-level tasks. Counting, reversing, spelling, character offsets, regex-like edits. Do them in code, not in the model.
- Numbers and identifiers. Anything you would be unhappy to see off by one digit should be validated after generation, or produced by a tool call.
What a real project would use instead
Nobody trains their own tokenizer for an application; you load the one the model was trained with, and it is not optional — a different tokenizer means a different meaning for every integer, and the model produces nonsense. The library is fast because the merges live in a compiled trie rather than the loop above, but the algorithm is the one you just wrote.
# The same three operations with the tokenizers a real project uses.
# NOT RUN HERE - tiktoken and transformers are not installed.
import tiktoken
enc = tiktoken.get_encoding("cl100k_base") # the tokenizer, not the model
ids = enc.encode("the cat sat on the mat") # -> [1820, 8415, 7731, 389, 279, 5634]
enc.decode(ids) # back to the text, exactly
len(enc.encode(text)) # what you are billed for
# The letter-counting problem, visible:
[enc.decode([i]) for i in enc.encode("strawberry")] # ['str', 'aw', 'berry']
[enc.decode([i]) for i in enc.encode("s t r a w b e r r y")] # one token per letter
# The leading-space trap, visible:
enc.encode("cat") # one id
enc.encode(" cat") # a DIFFERENT id
# The non-English cost, measured on your own traffic rather than assumed:
en = len(enc.encode("the cat sat on the mat"))
ja = len(enc.encode("猫がマットの上に座った"))
ratio = ja / en # budget with this, not with a rule of thumb
# For an open model, the tokenizer ships with the weights and must match them.
from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B")
tok("the cat sat on the mat")["input_ids"]
tok.apply_chat_template(messages, tokenize=False) # the special tokens matter tootiktoken and transformers are not installed in this environment, so this panel was never executed and the token ids in the comments are quoted from documentation rather than measured here. The algorithm behind every call in it was written out and asserted in plain Python above.
The mismatch bug. Counting tokens with one tokenizer and sending them to a model trained on another is a silent failure, not a loud one. Your length checks pass, your cost estimate is wrong, your truncation cuts in the wrong place, and the output is merely a bit worse than it should be. Pin the tokenizer to the model, in the same config, and assert the pairing in a test. When a model version changes, the tokenizer may change too.
18 Predicting the next word
A language model answers one question: given what came before, what comes next. You can build one with nothing but a dictionary and division — count which words follow which, then normalise. That model works, is genuinely useful, and fails in four specific ways. Those four failures are the reason the rest of this part of the course exists.
A bigram model is a table of counts
A bigram model looks one word back. Count every adjacent pair in a corpus, divide each row by its total, and you have a probability distribution over the next word for every word in the vocabulary. There is no training loop, no loss and no gradient — counting is the whole thing.
# A bigram language model. Count pairs, divide, and that is the whole model.
import random
CORPUS = """the cat sat on the mat
the cat ate the fish
the dog sat on the log
the dog ate the bone
the cat sat on the log
the dog ate the fish"""
START = "<s>"
END = "</s>"
def sentences(text):
"""Markers at both ends, so 'what starts a sentence' and 'when to stop' are
ordinary predictions rather than special cases."""
return [[START] + line.split() + [END] for line in text.strip().split("\n")]
def bigram_counts(rows):
counts = {}
for row in rows:
for a, b in zip(row, row[1:]):
counts.setdefault(a, {})
counts[a][b] = counts[a].get(b, 0) + 1
return counts
def to_probs(counts):
model = {}
for a, nexts in counts.items():
total = sum(nexts.values())
model[a] = {b: n / total for b, n in nexts.items()}
return model
def sample(dist, rng):
"""Inverse transform sampling, over a sorted key order so it is reproducible."""
r, running = rng.random(), 0.0
for word, p in sorted(dist.items()):
running += p
if r < running:
return word
return sorted(dist)[-1]
def generate(model, rng, limit=12):
word, out = START, []
while len(out) < limit:
word = sample(model[word], rng)
if word == END:
break
out.append(word)
return out
rows = sentences(CORPUS)
counts = bigram_counts(rows)
model = to_probs(counts)
print("what follows 'the', and how often:")
for word, n in sorted(counts["the"].items(), key=lambda kv: -kv[1]):
print(" the {:6} {:3} times p = {:.3f}".format(word, n, model["the"][word]))
print()
print("what follows 'sat':", dict(sorted(counts["sat"].items())))
print("what follows 'ate':", dict(sorted(counts["ate"].items())))
rng = random.Random(42)
print()
print("six sentences generated from those counts alone:")
for _ in range(6):
print(" ", " ".join(generate(model, rng)))
first_run = [generate(model, random.Random(42)) for _ in range(1)]
repeat = [generate(model, random.Random(42)) for _ in range(1)]
test("every row of the model is a probability distribution",
all(close(sum(d.values()), 1.0) for d in model.values()), True)
test("six different words follow 'the' in this corpus",
sorted(model["the"]), ["bone", "cat", "dog", "fish", "log", "mat"])
test("'on' is followed by 'the' every single time", close(model["on"]["the"], 1.0), True)
test("the same seed replays the same sentence", first_run, repeat)
test("the model never invents a word the corpus did not contain",
all(w in CORPUS.split() for row in first_run for w in row), True)
test("counting is the whole of training: no loss, no gradient, no epochs",
sum(sum(d.values()) for d in counts.values()), sum(len(r) - 1 for r in rows))What it printed, when the build ran it
what follows 'the', and how often:
the cat 3 times p = 0.250
the dog 3 times p = 0.250
the fish 2 times p = 0.167
the log 2 times p = 0.167
the mat 1 times p = 0.083
the bone 1 times p = 0.083
what follows 'sat': {'on': 3}
what follows 'ate': {'the': 3}
six sentences generated from those counts alone:
the bone
the fish
the cat sat on the dog ate the fish
the fish
the log
the cat sat on the cat sat on the fish
ok every row of the model is a probability distribution
ok six different words follow 'the' in this corpus
ok 'on' is followed by 'the' every single time
ok the same seed replays the same sentence
ok the model never invents a word the corpus did not contain
ok counting is the whole of training: no loss, no gradient, no epochs
Read the generated sentences honestly. Locally they are fine — every pair of adjacent words is a pair that really occurs. Globally some are nonsense, because the model has already forgotten everything except the last word. "the bone" is a complete sentence as far as this model is concerned. That is not a bug in the counting; it is the definition of a bigram.
Failure one: the table is almost all zeros
A vocabulary of V words has V squared possible pairs. A real corpus contains a tiny fraction of them, and every pair it missed is assigned probability zero — not "unlikely", zero. Multiply one zero through a sentence and the whole sentence becomes impossible.
# Sparsity. Most pairs of words never happened, and a count of zero is fatal.
import math
CORPUS = """the cat sat on the mat
the cat ate the fish
the dog sat on the log
the dog ate the bone
the cat sat on the log
the dog ate the fish"""
START = "<s>"
END = "</s>"
def sentences(text):
return [[START] + line.split() + [END] for line in text.strip().split("\n")]
def bigram_counts(rows):
counts = {}
for row in rows:
for a, b in zip(row, row[1:]):
counts.setdefault(a, {})
counts[a][b] = counts[a].get(b, 0) + 1
return counts
rows = sentences(CORPUS)
counts = bigram_counts(rows)
vocab = sorted(set(w for row in rows for w in row))
def p_bigram(a, b):
nexts = counts.get(a, {})
total = sum(nexts.values())
return 0.0 if total == 0 else nexts.get(b, 0) / total
def p_sentence(words):
row = [START] + words + [END]
p = 1.0
for a, b in zip(row, row[1:]):
p *= p_bigram(a, b)
return p
possible = len(vocab) * len(vocab)
seen = sum(len(n) for n in counts.values())
print("vocabulary :", vocab)
print("possible pairs :", possible)
print("pairs ever seen :", seen)
print("the table is : {:.1f}% zeros".format(100.0 * (possible - seen) / possible))
print()
print(" sentence probability")
for words in (["the", "cat", "sat", "on", "the", "mat"],
["the", "cat", "sat", "on", "the", "bone"],
["the", "fish", "ate", "the", "cat"],
["the", "cat", "sat", "on", "the", "cat"]):
print(" {:30} {}".format(" ".join(words), p_sentence(words)))
print()
print("Two failures in one table. 'on the bone' scores exactly as well as 'on the")
print("mat', because a bigram model has already forgotten the word 'on'. And 'the")
print("fish ate' scores zero, because that pair never happened, and one zero")
print("multiplied through the product kills the entire sentence.")
real = p_sentence(["the", "cat", "sat", "on", "the", "mat"])
odd = p_sentence(["the", "cat", "sat", "on", "the", "bone"])
unseen = p_sentence(["the", "fish", "ate", "the", "cat"])
try:
math.log(unseen)
logged = True
except ValueError as error:
logged = False
print()
print("math.log(0.0) ->", error, " <- so there is no loss to report either")
test("a sentence from the corpus has a real probability", real > 0.0, True)
test("a sentence containing one unseen pair is given probability zero",
close(unseen, 0.0), True)
test("the bigram model cannot tell the odd sentence from the real one",
close(real, odd), True)
test("because it only ever looks one word back", p_bigram("the", "bone") > 0.0, True)
test("almost seven eighths of the pair table is zeros",
close((possible - seen) / possible, 0.875), True)
test("and a zero probability has no logarithm, so the model cannot be scored",
logged, False)What it printed, when the build ran it
vocabulary : ['</s>', '<s>', 'ate', 'bone', 'cat', 'dog', 'fish', 'log', 'mat', 'on', 'sat', 'the'] possible pairs : 144 pairs ever seen : 18 the table is : 87.5% zeros sentence probability the cat sat on the mat 0.013888888888888888 the cat sat on the bone 0.013888888888888888 the fish ate the cat 0.0 the cat sat on the cat 0.0 Two failures in one table. 'on the bone' scores exactly as well as 'on the mat', because a bigram model has already forgotten the word 'on'. And 'the fish ate' scores zero, because that pair never happened, and one zero multiplied through the product kills the entire sentence. math.log(0.0) -> math domain error <- so there is no loss to report either ok a sentence from the corpus has a real probability ok a sentence containing one unseen pair is given probability zero ok the bigram model cannot tell the odd sentence from the real one ok because it only ever looks one word back ok almost seven eighths of the pair table is zeros ok and a zero probability has no logarithm, so the model cannot be scored
Sparsity gets worse, not better, as the corpus grows, because the vocabulary grows with it. Ten thousand words is a hundred million possible pairs and a billion possible triples, and no corpus on earth contains most of them. This is the wall that counting hits and the reason neural language models exist: an embedding lets "fish" and "cat" share what they have in common, so a pair never seen can still be scored. A count table cannot generalise by construction — every cell is independent of every other cell.
Failure two, patched: add-k smoothing
The cheapest fix is to lie a little. Pretend you saw every possible pair k
extra times before you started counting. With k = 1 it is called Laplace or
add-one smoothing. Nothing is impossible any more, and every row still sums to one.
# Add-k smoothing: pretend you saw every pair k extra times. Laplace is k = 1.
CORPUS = """the cat sat on the mat
the cat ate the fish
the dog sat on the log
the dog ate the bone
the cat sat on the log
the dog ate the fish"""
START = "<s>"
END = "</s>"
def sentences(text):
return [[START] + line.split() + [END] for line in text.strip().split("\n")]
def bigram_counts(rows):
counts = {}
for row in rows:
for a, b in zip(row, row[1:]):
counts.setdefault(a, {})
counts[a][b] = counts[a].get(b, 0) + 1
return counts
rows = sentences(CORPUS)
counts = bigram_counts(rows)
vocab = sorted(set(w for row in rows for w in row))
V = len(vocab)
def p_add_k(a, b, k):
"""Every one of the V possible next words gets k free sightings."""
if k <= 0.0:
raise ValueError("k must be positive, or nothing is smoothed")
nexts = counts.get(a, {})
return (nexts.get(b, 0) + k) / (sum(nexts.values()) + k * V)
print("vocabulary size V =", V)
print()
print(" k p(cat|the) p(fish|the) p(sat|the) sums over V")
for k in (0.001, 0.01, 0.1, 1.0, 100.0):
row = [p_add_k("the", w, k) for w in vocab]
print(" {:7} {:10.4f} {:11.4f} {:10.4f} {:.6f}".format(
k, p_add_k("the", "cat", k), p_add_k("the", "fish", k),
p_add_k("the", "sat", k), sum(row)))
print()
print("'sat' never follows 'the' in this corpus. Unsmoothed it is impossible.")
print("At k = 1 it is merely unlikely, which is a claim you can take a log of.")
print("At k = 100 the counts have been drowned and everything is nearly uniform.")
unsmoothed = counts["the"].get("sat", 0) / sum(counts["the"].values())
laplace_unseen = p_add_k("the", "sat", 1.0)
laplace_seen = p_add_k("the", "cat", 1.0)
raw_seen = counts["the"]["cat"] / sum(counts["the"].values())
flat = [p_add_k("the", w, 100.0) for w in vocab]
try:
p_add_k("the", "cat", 0.0)
refused = False
except ValueError as error:
refused = True
print()
print("k = 0 is refused:", error)
test("the smoothed row is still a distribution over the whole vocabulary",
close(sum(p_add_k("the", w, 1.0) for w in vocab), 1.0), True)
test("a pair that never happened now has a probability above zero",
laplace_unseen > 0.0, True)
test("and with k = 1 it is exactly one over the count plus the vocabulary size",
close(laplace_unseen, 1.0 / (12 + 12)), True)
test("the mass had to come from somewhere, so a pair that did happen drops",
laplace_seen < raw_seen, True)
test("unsmoothed, that same unseen pair was flatly impossible",
close(unsmoothed, 0.0), True)
test("a very large k drowns the evidence and returns a uniform distribution",
all(close(x, 1.0 / V, 5e-3) for x in flat), True)
test("k must be positive, or nothing has been smoothed at all", refused, True)What it printed, when the build ran it
vocabulary size V = 12
k p(cat|the) p(fish|the) p(sat|the) sums over V
0.001 0.2498 0.1666 0.0001 1.000000
0.01 0.2483 0.1658 0.0008 1.000000
0.1 0.2348 0.1591 0.0076 1.000000
1.0 0.1667 0.1250 0.0417 1.000000
100.0 0.0850 0.0842 0.0825 1.000000
'sat' never follows 'the' in this corpus. Unsmoothed it is impossible.
At k = 1 it is merely unlikely, which is a claim you can take a log of.
At k = 100 the counts have been drowned and everything is nearly uniform.
k = 0 is refused: k must be positive, or nothing is smoothed
ok the smoothed row is still a distribution over the whole vocabulary
ok a pair that never happened now has a probability above zero
ok and with k = 1 it is exactly one over the count plus the vocabulary size
ok the mass had to come from somewhere, so a pair that did happen drops
ok unsmoothed, that same unseen pair was flatly impossible
ok a very large k drowns the evidence and returns a uniform distribution
ok k must be positive, or nothing has been smoothed at all
Notice the trade in the first two columns. p(cat|the) falls from 0.25 to 0.17
as k goes to 1, because the probability given to pairs that never happened has to
be taken from pairs that did. Smoothing does not add information. It spends accuracy on the
words you saw to buy coverage of the words you did not, and k is the exchange
rate. Choosing it by hand on the training set defeats the purpose; you choose it on a
held-out split, which is the next section.
Add-k is the version worth writing once and then never using. Real count-based models
use backoff or interpolation: if the trigram
(a, b, c) was never seen, fall back to the bigram (b, c), and
below that to the unigram frequency of c. That is smarter because it spends
the leftover mass in proportion to what the shorter context already knows, rather than
spreading it uniformly over words that were never plausible. Kneser-Ney is the refined
version of the same idea, and was the state of the art until neural models arrived.
Perplexity: one number for how surprised the model was
Take the probability the model gave to each word that actually occurred. Take the log of each, average them, negate, exponentiate. The result has a reading in plain English: the effective number of equally likely words the model was choosing between at each step. Lower is better and the floor is one.
# Perplexity. One number for "how surprised was the model, on average".
import math
CORPUS = """the cat sat on the mat
the cat ate the fish
the dog sat on the log
the dog ate the bone
the cat sat on the log
the dog ate the fish"""
HELD_OUT = """the dog sat on the mat
the cat ate the bone
the fish sat on the log"""
START = "<s>"
END = "</s>"
def sentences(text):
return [[START] + line.split() + [END] for line in text.strip().split("\n")]
def bigram_counts(rows):
counts = {}
for row in rows:
for a, b in zip(row, row[1:]):
counts.setdefault(a, {})
counts[a][b] = counts[a].get(b, 0) + 1
return counts
rows = sentences(CORPUS)
held = sentences(HELD_OUT)
counts = bigram_counts(rows)
vocab = sorted(set(w for row in rows for w in row))
V = len(vocab)
def p_add_k(a, b, k=1.0):
nexts = counts.get(a, {})
return (nexts.get(b, 0) + k) / (sum(nexts.values()) + k * V)
def perplexity(rows, prob):
"""exp of the average negative log probability, per predicted word."""
total_log, n = 0.0, 0
for row in rows:
for a, b in zip(row, row[1:]):
p = prob(a, b)
if p <= 0.0:
return float("inf") # one impossible word and the score is infinite
total_log += math.log(p)
n += 1
return math.exp(-total_log / n)
def uniform(a, b):
return 1.0 / V
def oracle(a, b):
return 1.0
def bigram(a, b):
return p_add_k(a, b, 1.0)
def raw_bigram(a, b):
nexts = counts.get(a, {})
total = sum(nexts.values())
return 0.0 if total == 0 else nexts.get(b, 0) / total
print("vocabulary size V =", V)
print()
print(" model on the corpus on held-out text")
for name, prob in (("uniform guess, 1/V every time", uniform),
("smoothed bigram, k = 1", bigram),
("an oracle that is always certain", oracle)):
print(" {:33} {:14.3f} {:.3f}".format(
name, perplexity(rows, prob), perplexity(held, prob)))
print(" {:33} {:14.3f} {}".format(
"unsmoothed bigram", perplexity(rows, raw_bigram), perplexity(held, raw_bigram)))
print()
print("Read perplexity as a number of choices. 12 means the model was worth no")
print("more than picking uniformly from 12 words. 5 means it had narrowed the")
print("field to about five. 1 means it knew. The last row is the point of")
print("smoothing: one pair the corpus never held, and the score is infinite.")
train_ppl = perplexity(rows, bigram)
held_ppl = perplexity(held, bigram)
test("a uniform model over twelve words has perplexity of exactly twelve",
close(perplexity(rows, uniform), 12.0), True)
test("a model that is always certain and always right has perplexity one",
close(perplexity(rows, oracle), 1.0), True)
test("the bigram model beats guessing on the text it was counted from",
train_ppl < 12.0, True)
test("the same model is more surprised by text it has not seen", held_ppl > train_ppl, True)
test("smoothing keeps the held-out score finite", math.isinf(held_ppl), False)
test("without smoothing one unseen pair makes it infinite",
math.isinf(perplexity(held, raw_bigram)), True)
test("perplexity never drops below one, because probability never exceeds one",
min(train_ppl, held_ppl) >= 1.0, True)What it printed, when the build ran it
vocabulary size V = 12 model on the corpus on held-out text uniform guess, 1/V every time 12.000 12.000 smoothed bigram, k = 1 4.899 5.503 an oracle that is always certain 1.000 1.000 unsmoothed bigram 1.863 inf Read perplexity as a number of choices. 12 means the model was worth no more than picking uniformly from 12 words. 5 means it had narrowed the field to about five. 1 means it knew. The last row is the point of smoothing: one pair the corpus never held, and the score is infinite. ok a uniform model over twelve words has perplexity of exactly twelve ok a model that is always certain and always right has perplexity one ok the bigram model beats guessing on the text it was counted from ok the same model is more surprised by text it has not seen ok smoothing keeps the held-out score finite ok without smoothing one unseen pair makes it infinite ok perplexity never drops below one, because probability never exceeds one
The uniform row is the reason perplexity is worth learning: a uniform model over a vocabulary of V has perplexity exactly V, so the number is directly comparable to "how many words could it have been". It is also why a perplexity figure quoted without its test set and its tokenizer is meaningless. Perplexity is per token, so a model with a bigger vocabulary predicts fewer, harder tokens per sentence and gets a different number for identical text. Two models are only comparable on perplexity if they are scored on the same text with the same tokenizer.
Failure three: more context means memorising
The obvious fix for a bigram forgetting too much is to look further back. A trigram conditions on two words instead of one. On a small corpus that does not produce a better model; it produces a copy of the corpus.
# A trigram model. More context, and on a corpus this small it simply recites.
import math
import random
CORPUS = """the cat sat on the mat
a dog barked at the moon
my friend wrote a long letter
the rain fell all afternoon
she left before the film ended
we ate bread and cheese"""
HELD_OUT = """the dog sat on the moon
a cat barked at the mat"""
START = "<s>"
END = "</s>"
def sentences(text, back):
"""Pad with as many start markers as the context needs, so the first real
word is predicted from a context that exists."""
return [[START] * back + line.split() + [END] for line in text.strip().split("\n")]
def ngram_counts(rows, order):
"""order 2 is a bigram model, 3 a trigram. The context is order - 1 words."""
counts = {}
back = order - 1
for row in rows:
for i in range(back, len(row)):
context = tuple(row[i - back:i])
counts.setdefault(context, {})
counts[context][row[i]] = counts[context].get(row[i], 0) + 1
return counts
bi_rows, bi_held = sentences(CORPUS, 1), sentences(HELD_OUT, 1)
tri_rows, tri_held = sentences(CORPUS, 2), sentences(HELD_OUT, 2)
V = len(set(w for row in tri_rows for w in row))
bi = ngram_counts(bi_rows, 2)
tri = ngram_counts(tri_rows, 3)
def probability(counts, context, word, k):
nexts = counts.get(context, {})
total = sum(nexts.values())
if k <= 0.0:
return 0.0 if total == 0 else nexts.get(word, 0) / total
return (nexts.get(word, 0) + k) / (total + k * V)
def perplexity(rows, counts, order, k):
back = order - 1
total, n = 0.0, 0
for row in rows:
for i in range(back, len(row)):
p = probability(counts, tuple(row[i - back:i]), row[i], k)
if p <= 0.0:
return float("inf")
total += math.log(p)
n += 1
return math.exp(-total / n)
def generate(counts, rng, order, limit=10):
back = order - 1
context, out = tuple([START] * back), []
while len(out) < limit:
nexts = counts[context]
total = sum(nexts.values())
r, running, pick = rng.random(), 0.0, sorted(nexts)[-1]
for word in sorted(nexts):
running += nexts[word] / total
if r < running:
pick = word
break
if pick == END:
break
out.append(pick)
context = tuple(list(context)[1:] + [pick])
return " ".join(out)
forced = sum(1 for nexts in tri.values() if len(nexts) == 1)
print("trigram contexts :", len(tri))
print("contexts with only one option:", forced)
print("so at {:.0f}% of positions the model has no decision left to make."
.format(100.0 * forced / len(tri)))
rng = random.Random(42)
lines = [generate(tri, rng, 3) for _ in range(20)]
corpus_lines = set(CORPUS.strip().split("\n"))
print()
print("twenty samples from the trigram gave", len(set(lines)), "distinct sentences:")
for line in sorted(set(lines)):
print(" {:32} {}".format(line, "verbatim from the corpus" if line in corpus_lines else "NEW"))
print()
print("the bigram model on the same corpus, for contrast:")
for line in [generate(bi, random.Random(42), 2) for _ in range(1)] + \
[generate(bi, random.Random(7), 2) for _ in range(1)] + \
[generate(bi, random.Random(11), 2) for _ in range(1)]:
print(" {:34} {}".format(line, "verbatim" if line in corpus_lines else "new, and often nonsense"))
print()
print(" model on the corpus on held-out text")
print(" {:26} {:13.3f} {:.3f}".format(
"bigram, k = 1", perplexity(bi_rows, bi, 2, 1.0), perplexity(bi_held, bi, 2, 1.0)))
print(" {:26} {:13.3f} {:.3f}".format(
"trigram, k = 1", perplexity(tri_rows, tri, 3, 1.0), perplexity(tri_held, tri, 3, 1.0)))
print(" {:26} {:13.3f} {}".format(
"bigram, no smoothing", perplexity(bi_rows, bi, 2, 0.0), perplexity(bi_held, bi, 2, 0.0)))
print(" {:26} {:13.3f} {}".format(
"trigram, no smoothing", perplexity(tri_rows, tri, 3, 0.0), perplexity(tri_held, tri, 3, 0.0)))
train_tri = perplexity(tri_rows, tri, 3, 0.0)
train_bi = perplexity(bi_rows, bi, 2, 0.0)
held_tri = perplexity(tri_held, tri, 3, 1.0)
held_bi = perplexity(bi_held, bi, 2, 1.0)
test("almost every trigram context has exactly one continuation",
forced / len(tri) > 0.9, True)
test("so twenty samples produce nothing the corpus did not already contain",
all(line in corpus_lines for line in lines), True)
test("unsmoothed, the trigram is barely surprised by its own training text",
train_tri < 1.4, True)
test("which beats the bigram on training text, and means nothing at all",
train_tri < train_bi, True)
test("on held-out text the trigram is the worse model of the two",
held_tri > held_bi, True)
test("and unsmoothed it is infinitely worse",
math.isinf(perplexity(tri_held, tri, 3, 0.0)), True)What it printed, when the build ran it
trigram contexts : 34
contexts with only one option: 32
so at 94% of positions the model has no decision left to make.
twenty samples from the trigram gave 6 distinct sentences:
a dog barked at the moon verbatim from the corpus
my friend wrote a long letter verbatim from the corpus
she left before the film ended verbatim from the corpus
the cat sat on the mat verbatim from the corpus
the rain fell all afternoon verbatim from the corpus
we ate bread and cheese verbatim from the corpus
the bigram model on the same corpus, for contrast:
the cat sat on the moon new, and often nonsense
my friend wrote a long letter verbatim
she left before the mat new, and often nonsense
model on the corpus on held-out text
bigram, k = 1 16.287 19.883
trigram, k = 1 16.049 23.420
bigram, no smoothing 1.600 inf
trigram, no smoothing 1.308 inf
ok almost every trigram context has exactly one continuation
ok so twenty samples produce nothing the corpus did not already contain
ok unsmoothed, the trigram is barely surprised by its own training text
ok which beats the bigram on training text, and means nothing at all
ok on held-out text the trigram is the worse model of the two
ok and unsmoothed it is infinitely worse
This is the whole argument for a held-out split, in one table. The unsmoothed trigram scores 1.3 on the corpus it was built from and infinity on two sentences it has not seen. Nothing about the training number was a lie — it was computed correctly — it simply measured recall of the training set rather than any ability to model language. The modern version of this mistake is an evaluation set that overlaps the training data, or a benchmark whose answers are on the web. It looks like a result and it is a measurement of nothing.
The k = 1 rows carry a second, quieter lesson. Add-one smoothing over a
thirty-word vocabulary makes both models look terrible on training text, because most contexts
were seen once or twice and the fake counts outweigh the real ones. Smoothing strong enough to
cover the zeros is strong enough to bury the signal. That tension — coverage against sharpness
— is what backoff schemes exist to manage, and what a neural model dissolves by sharing
structure between words instead of patching each cell independently.
- Count, then normalise. Every language model is estimating the same conditional distribution. The n-gram does it with division; a transformer does it with gradient descent, and the target is identical.
- Zeros are the enemy. Any model that assigns exactly zero to real text cannot be scored, cannot be trained by a log loss, and cannot be improved.
- Perplexity is per token. Compare it only on the same text with the same tokenizer, or you are comparing nothing.
- Held-out or it did not happen. A model with enough capacity will memorise, and the training score will look excellent while it does.
What a real project would use instead
The count table becomes a learned embedding table, and the division becomes gradient descent on the same objective — predict the next token, score it by log probability. Written in torch, the neural bigram below is about as long as the counting version above, and the thing it gains is exactly what counting cannot do: an unseen pair still gets a sensible score, because similar words have ended up with similar vectors.
# The same objective, learned instead of counted.
# NOT RUN HERE - torch is not installed.
import torch
import torch.nn.functional as F
V = 12 # vocabulary size
# The count table above was V x V numbers you divided. This is V x V numbers you
# learn, which is the entire difference in the setup.
logits_table = torch.zeros((V, V), requires_grad=True)
context = torch.tensor([0, 3, 5]) # previous token ids
target = torch.tensor([3, 5, 7]) # the token that actually came next
for step in range(200):
logits = logits_table[context] # one row of scores per context token
loss = F.cross_entropy(logits, target) # mean negative log probability
loss.backward()
with torch.no_grad():
logits_table -= 0.1 * logits_table.grad
logits_table.grad.zero_()
# Perplexity is exp of that same loss. It is not a separate metric.
perplexity = torch.exp(loss)
# And the reason anyone bothers: replace the row lookup with an embedding plus a
# hidden layer, and words that behave alike share parameters, so a pair the corpus
# never contained still gets a sensible probability. A count table cannot do that,
# because every cell in it is independent of every other cell.
model = torch.nn.Sequential(
torch.nn.Embedding(V, 16),
torch.nn.Flatten(),
torch.nn.Linear(16, 64),
torch.nn.Tanh(),
torch.nn.Linear(64, V),
)torch is not installed in this environment, so this panel was never executed and no number in it was measured here. The objective it minimises is the same average negative log probability that the perplexity function above computes in plain Python, and that one did run.
Everything a transformer does sits on top of what you just built. The objective is unchanged: put a distribution over the next token and make the real token likely. What changes is the context. A bigram sees one word, a trigram two, an n-gram n minus one — and each extra word multiplies the table size by the vocabulary. Attention, in lesson 20, gets arbitrary context without that explosion, by learning what to look at instead of enumerating every possibility.
19 Temperature, top-k and top-p
A model hands you one number per token and then does nothing else. Choosing a token from those numbers is called decoding, and it is entirely outside the model: no weights, no training, just arithmetic you control at request time. Four settings cover almost every API you will meet, and each one is a way of deciding how much of the distribution to take seriously.
Greedy and temperature, the two ends
Greedy takes the highest-scoring token, every time, with no randomness at all. Temperature divides the logits before the softmax, which sharpens the distribution below 1 and flattens it above 1. Neither one throws any token away.
# The two simplest decoders, both reading the same logits.
import math
import random
TOKENS = ("cat", "dog", "bird", "fish", "quark")
LOGITS = [3.0, 1.5, 0.5, 0.0, -2.0]
def softmax(logits, temperature=1.0):
if temperature <= 0.0:
raise ValueError("temperature must be positive; zero is argmax, not sampling")
scaled = [z / temperature for z in logits]
shift = max(scaled) # the overflow guard from lesson 6
exps = [math.exp(z - shift) for z in scaled]
total = sum(exps)
return [e / total for e in exps]
def greedy(probs):
"""No dice. The highest probability, every time."""
return probs.index(max(probs))
def sample(probs, rng):
r, running = rng.random(), 0.0
for i, p in enumerate(probs):
running += p
if r < running:
return i
return len(probs) - 1
def draw(probs, seed, n):
rng = random.Random(seed) # our own generator, never random.seed()
return [TOKENS[sample(probs, rng)] for _ in range(n)]
print("logits:", dict(zip(TOKENS, LOGITS)))
print()
print(" temperature " + "".join("{:>8}".format(t) for t in TOKENS))
for t in (0.2, 0.7, 1.0, 1.5, 3.0):
print(" {:11} ".format(t) + "".join("{:8.4f}".format(x) for x in softmax(LOGITS, t)))
base = softmax(LOGITS)
print()
print("greedy, ten times :", [TOKENS[greedy(base)] for _ in range(10)])
print("sampled at T=1.0 :", draw(base, 42, 10))
print("sampled at T=0.2 :", draw(softmax(LOGITS, 0.2), 42, 10))
print("sampled at T=3.0 :", draw(softmax(LOGITS, 3.0), 42, 10))
hot = softmax(LOGITS, 3.0)
cold = softmax(LOGITS, 0.2)
test("softmax gives a distribution at every temperature",
all(close(sum(softmax(LOGITS, t)), 1.0) for t in (0.2, 1.0, 3.0)), True)
test("greedy is deterministic, so ten runs give one answer",
len(set(TOKENS[greedy(base)] for _ in range(10))), 1)
test("a cold temperature puts nearly all the mass on that same token", cold[0] > 0.99, True)
test("a hot temperature makes even the least likely token reachable", hot[4] > 0.03, True)
test("temperature never reorders the tokens",
[greedy(softmax(LOGITS, t)) for t in (0.2, 1.0, 3.0)], [0, 0, 0])
test("the same seed replays the same draws", draw(base, 42, 10), draw(base, 42, 10))What it printed, when the build ran it
logits: {'cat': 3.0, 'dog': 1.5, 'bird': 0.5, 'fish': 0.0, 'quark': -2.0}
temperature cat dog bird fish quark
0.2 0.9994 0.0006 0.0000 0.0000 0.0000
0.7 0.8621 0.1011 0.0242 0.0119 0.0007
1.0 0.7344 0.1639 0.0603 0.0366 0.0049
1.5 0.5788 0.2129 0.1093 0.0783 0.0206
3.0 0.3849 0.2335 0.1673 0.1416 0.0727
greedy, ten times : ['cat', 'cat', 'cat', 'cat', 'cat', 'cat', 'cat', 'cat', 'cat', 'cat']
sampled at T=1.0 : ['cat', 'cat', 'cat', 'cat', 'dog', 'cat', 'dog', 'cat', 'cat', 'cat']
sampled at T=0.2 : ['cat', 'cat', 'cat', 'cat', 'cat', 'cat', 'cat', 'cat', 'cat', 'cat']
sampled at T=3.0 : ['bird', 'cat', 'cat', 'cat', 'bird', 'bird', 'fish', 'cat', 'dog', 'cat']
ok softmax gives a distribution at every temperature
ok greedy is deterministic, so ten runs give one answer
ok a cold temperature puts nearly all the mass on that same token
ok a hot temperature makes even the least likely token reachable
ok temperature never reorders the tokens
ok the same seed replays the same draws
Temperature is a volume knob, not a steering wheel. It changes how much of the
distribution you are willing to visit; it never changes the ranking. That matters because
"turn up the temperature to get a more creative answer" is only half true — at
T = 3.0 the model still says "cat" most often, it just now sometimes says
"quark". The problem with temperature alone is exactly that: a high setting leaves genuinely
bad tokens reachable, and a long generation will eventually reach one.
top-k, and where a fixed k stops making sense
top-k keeps the k highest probabilities, sets the rest to zero, and renormalises so the survivors sum to one again. It puts a hard floor under quality: whatever the temperature, the 40,000th-ranked token cannot be chosen. The cost is that k is a constant and the distribution is not.
# top-k: keep the k best, throw the rest away, renormalise. Then look at
# where a fixed k stops making sense.
import math
PEAKED = [6.0, 2.0, 1.0, 0.5, 0.0, -0.5, -1.0, -1.5, -2.0, -2.5]
FLAT = [1.00, 0.99, 0.98, 0.97, 0.96, 0.95, 0.94, 0.93, 0.92, 0.91]
def softmax(logits):
shift = max(logits)
exps = [math.exp(z - shift) for z in logits]
total = sum(exps)
return [e / total for e in exps]
def top_k(probs, k):
"""Keep the k highest probabilities, zero the rest, renormalise to sum to 1."""
if k < 1 or k > len(probs):
raise ValueError("k must be between 1 and the vocabulary size")
ranked = sorted(range(len(probs)), key=lambda i: -probs[i])
keep = set(ranked[:k])
kept_mass = sum(probs[i] for i in keep)
return [probs[i] / kept_mass if i in keep else 0.0 for i in range(len(probs))]
def alive(dist):
return sum(1 for x in dist if x > 0.0)
peaked = softmax(PEAKED)
flat = softmax(FLAT)
print("a confident step. one token holds almost all the mass:")
print(" ", ["{:.4f}".format(x) for x in peaked])
print("an unsure step. ten tokens, all about as likely as each other:")
print(" ", ["{:.4f}".format(x) for x in flat])
print()
print(" k tokens left mass kept before renormalising which step")
for k in (1, 3, 5, 10):
print(" {:3} {:11} {:29.4f} confident".format(
k, alive(top_k(peaked, k)), sum(sorted(peaked, reverse=True)[:k])))
for k in (1, 3, 5, 10):
print(" {:3} {:11} {:29.4f} unsure".format(
k, alive(top_k(flat, k)), sum(sorted(flat, reverse=True)[:k])))
print()
print("On the confident step, k = 3 throws away 1% of the mass and changes")
print("nothing. On the unsure step, the same k = 3 throws away 69% of a")
print("distribution that genuinely had no favourite.")
ranked = sorted(range(10), key=lambda i: -flat[i])
inside, outside = ranked[2], ranked[3]
print()
print("and the cut is arbitrary. on the unsure step, with k = 3:")
print(" last token kept : p = {:.6f}".format(flat[inside]))
print(" first token cut : p = {:.6f}".format(flat[outside]))
print(" difference : {:.6f}".format(flat[inside] - flat[outside]))
try:
top_k(flat, 0)
refused = False
except ValueError as error:
refused = True
print()
print("k = 0 is refused:", error)
peaked_kept = sum(sorted(peaked, reverse=True)[:3])
flat_kept = sum(sorted(flat, reverse=True)[:3])
test("top-k leaves exactly k tokens alive", alive(top_k(flat, 3)), 3)
test("and the survivors still form a distribution", close(sum(top_k(flat, 3)), 1.0), True)
test("k = 1 is greedy decoding with extra steps", alive(top_k(peaked, 1)), 1)
test("on a confident step, k = 3 keeps nearly all the mass", peaked_kept > 0.98, True)
test("on an unsure step, the same k throws most of it away", flat_kept < 0.35, True)
test("and the token just outside the cut was as likely as the one just inside",
close(flat[inside], flat[outside], 0.002), True)
test("a k outside 1 to V is refused rather than silently clamped", refused, True)What it printed, when the build ran it
a confident step. one token holds almost all the mass:
['0.9661', '0.0177', '0.0065', '0.0039', '0.0024', '0.0015', '0.0009', '0.0005', '0.0003', '0.0002']
an unsure step. ten tokens, all about as likely as each other:
['0.1046', '0.1035', '0.1025', '0.1015', '0.1005', '0.0995', '0.0985', '0.0975', '0.0965', '0.0956']
k tokens left mass kept before renormalising which step
1 1 0.9661 confident
3 3 0.9903 confident
5 5 0.9966 confident
10 10 1.0000 confident
1 1 0.1046 unsure
3 3 0.3106 unsure
5 5 0.5125 unsure
10 10 1.0000 unsure
On the confident step, k = 3 throws away 1% of the mass and changes
nothing. On the unsure step, the same k = 3 throws away 69% of a
distribution that genuinely had no favourite.
and the cut is arbitrary. on the unsure step, with k = 3:
last token kept : p = 0.102489
first token cut : p = 0.101469
difference : 0.001020
k = 0 is refused: k must be between 1 and the vocabulary size
ok top-k leaves exactly k tokens alive
ok and the survivors still form a distribution
ok k = 1 is greedy decoding with extra steps
ok on a confident step, k = 3 keeps nearly all the mass
ok on an unsure step, the same k throws most of it away
ok and the token just outside the cut was as likely as the one just inside
ok a k outside 1 to V is refused rather than silently clamped
Read the last three lines of that output again. With k = 3 on the unsure
step, the token ranked third survives with probability 0.1025 and the token ranked fourth
is deleted with probability 0.1015. A gap of one part in a thousand decided whether a token
was possible or impossible. That is not a judgement about language, it is an artefact of
picking a constant. It is also why k is the setting people most often get
wrong: small values strangle the model exactly when it needed room, and large values do
nothing at all when it did not.
top-p, which measures the distribution instead of assuming it
top-p, also called nucleus sampling, fixes a mass rather than a count. Sort the tokens by probability, walk down the list adding them up, and stop as soon as the running total reaches p. The number of tokens kept is whatever that took, so it is narrow where the model is confident and wide where it is not.
# top-p, or nucleus sampling. Choose the mass to keep, and let the width follow.
import math
PEAKED = [6.0, 2.0, 1.0, 0.5, 0.0, -0.5, -1.0, -1.5, -2.0, -2.5]
FLAT = [1.00, 0.99, 0.98, 0.97, 0.96, 0.95, 0.94, 0.93, 0.92, 0.91]
def softmax(logits):
shift = max(logits)
exps = [math.exp(z - shift) for z in logits]
total = sum(exps)
return [e / total for e in exps]
def top_k(probs, k):
ranked = sorted(range(len(probs)), key=lambda i: -probs[i])
keep = set(ranked[:k])
mass = sum(probs[i] for i in keep)
return [probs[i] / mass if i in keep else 0.0 for i in range(len(probs))]
def top_p(probs, p):
"""Keep the smallest set of tokens whose probabilities reach p, then renormalise."""
if not 0.0 < p <= 1.0:
raise ValueError("p must be above 0 and at most 1")
ranked = sorted(range(len(probs)), key=lambda i: -probs[i])
keep, running = [], 0.0
for i in ranked:
keep.append(i)
running += probs[i]
if running >= p:
break # the token that crosses p is kept, not dropped
mass = sum(probs[i] for i in keep)
kept = set(keep)
return [probs[i] / mass if i in kept else 0.0 for i in range(len(probs))]
def alive(dist):
return sum(1 for x in dist if x > 0.0)
def entropy(dist):
return -sum(x * math.log(x) for x in dist if x > 0.0)
peaked = softmax(PEAKED)
flat = softmax(FLAT)
print(" p tokens kept, confident step tokens kept, unsure step")
for p in (0.5, 0.8, 0.9, 0.95, 1.0):
print(" {:5} {:28} {}".format(p, alive(top_p(peaked, p)), alive(top_p(flat, p))))
print()
print(" the same two steps under a fixed k")
print(" k tokens kept, confident step tokens kept, unsure step")
for k in (1, 3, 5, 10):
print(" {:5} {:28} {}".format(k, alive(top_k(peaked, k)), alive(top_k(flat, k))))
print()
print("That is the whole argument. top-p read the shape of the distribution and")
print("kept one token where the model was sure and nine where it was not.")
print("top-k kept three either way, because three is what you told it.")
print()
print("entropy, which is what top-p is implicitly reacting to:")
print(" confident step: {:.4f} nats".format(entropy(peaked)))
print(" unsure step : {:.4f} nats".format(entropy(flat)))
try:
top_p(flat, 0.0)
refused = False
except ValueError as error:
refused = True
print()
print("p = 0 is refused:", error)
test("top-p leaves a distribution behind", close(sum(top_p(flat, 0.9)), 1.0), True)
test("on a confident step, p = 0.9 needs only one token", alive(top_p(peaked, 0.9)), 1)
test("on an unsure step, the same p needs nine", alive(top_p(flat, 0.9)), 9)
test("while a fixed k keeps the same number for both",
alive(top_k(peaked, 3)), alive(top_k(flat, 3)))
test("p = 1.0 keeps the whole vocabulary", alive(top_p(flat, 1.0)), 10)
test("a tiny p collapses to greedy, because one token already covers it",
alive(top_p(peaked, 0.01)), 1)
test("p outside its range is refused", refused, True)What it printed, when the build ran it
p tokens kept, confident step tokens kept, unsure step
0.5 1 5
0.8 1 8
0.9 1 9
0.95 1 10
1.0 10 10
the same two steps under a fixed k
k tokens kept, confident step tokens kept, unsure step
1 1 1
3 3 3
5 5 5
10 10 10
That is the whole argument. top-p read the shape of the distribution and
kept one token where the model was sure and nine where it was not.
top-k kept three either way, because three is what you told it.
entropy, which is what top-p is implicitly reacting to:
confident step: 0.1978 nats
unsure step : 2.3022 nats
p = 0 is refused: p must be above 0 and at most 1
ok top-p leaves a distribution behind
ok on a confident step, p = 0.9 needs only one token
ok on an unsure step, the same p needs nine
ok while a fixed k keeps the same number for both
ok p = 1.0 keeps the whole vocabulary
ok a tiny p collapses to greedy, because one token already covers it
ok p outside its range is refused
top-p is not strictly better, and the failure worth knowing is the mirror image of
top-k's. When the model is confident, top-p with a high p can still admit a long tail: if
the best token has probability 0.55, p = 0.95 keeps everything down to 0.95 of
the mass, which on a 50,000-token vocabulary can be hundreds of tokens. That is why real
APIs expose both and why the common recipe is a moderate p with a
k behind it as a hard ceiling. The two filters are cheap; use them
together.
Repetition, and why there is a setting for it
Greedy decoding loops. Not sometimes — structurally. If the best token after A is B and the best token after B is A, greedy has no way out, and the loop is the correct answer to the question greedy is asking. The repetition penalty is a patch applied to the logits of tokens that have already appeared.
# Repetition, and the blunt instrument that exists to stop it.
SCORES = {
"<s>": {"the": 4.0, "a": 2.0},
"the": {"cat": 4.0, "mat": 3.5, "dog": 2.0, "</s>": -1.0},
"cat": {"the": 4.0, "sat": 3.0, "</s>": 0.0},
"sat": {"on": 4.0, "</s>": 0.0},
"on": {"the": 4.0, "</s>": 0.0},
"mat": {"</s>": 4.0, "the": 1.0},
"dog": {"the": 3.0, "sat": 2.0},
"a": {"cat": 3.0, "dog": 2.0},
}
def penalise(scores, used, penalty):
"""The rule a real library uses: divide a positive logit, multiply a negative
one. Dividing a negative number would make it larger, which is backwards."""
if penalty < 1.0:
raise ValueError("a penalty below 1 rewards repetition instead of discouraging it")
out = {}
for token, z in scores.items():
if token in used:
out[token] = z / penalty if z > 0.0 else z * penalty
else:
out[token] = z
return out
def greedy_walk(penalty=1.0, limit=12):
token, out, used = "<s>", [], set()
while len(out) < limit:
scores = penalise(SCORES[token], used, penalty)
token = max(sorted(scores), key=lambda t: scores[t])
if token == "</s>":
break
out.append(token)
used.add(token)
return out
loop = greedy_walk(1.0)
fixed = greedy_walk(1.5)
print("greedy, no penalty :", " ".join(loop))
print("greedy, penalty 1.5 :", " ".join(fixed))
print()
print("Nothing is broken. Greedy takes the best next token, the best token after")
print("'cat' is 'the', and the best token after 'the' is 'cat'. The loop is the")
print("correct answer to the question greedy decoding is asking.")
print()
print("what the penalty does to one step. after 'the cat', the scores from 'cat':")
for name, penalty in (("no penalty", 1.0), ("penalty 1.5", 1.5), ("penalty 3.0", 3.0)):
scores = penalise(SCORES["cat"], {"the", "cat"}, penalty)
print(" {:12}".format(name), {k: round(v, 3) for k, v in sorted(scores.items())})
print()
print("The end marker at 0.0 does not move: the rule only touches tokens already")
print("used. And 'the' is being punished for being a word English needs to repeat.")
try:
greedy_walk(0.5)
refused = False
except ValueError as error:
refused = True
print()
print("a penalty below 1 is refused:", error)
untouched = penalise(SCORES["the"], {"the", "cat"}, 2.0)
test("greedy with no penalty runs into a two-token loop", loop, loop[:2] * 6)
test("and never reaches the end marker, it just hits the length limit", len(loop), 12)
test("a penalty of 1.5 breaks the loop",
fixed, ["the", "cat", "sat", "on", "the", "mat"])
test("and lets the sentence finish on its own", len(fixed) < 12, True)
test("a penalty of 1.0 changes nothing at all", greedy_walk(1.0), loop)
test("the penalty leaves tokens that have not been used alone",
close(untouched["mat"], 3.5), True)
test("but it does punish a word that English needs to repeat",
untouched["cat"] < SCORES["the"]["cat"], True)
test("a penalty below one is refused, because it would reward repetition", refused, True)What it printed, when the build ran it
greedy, no penalty : the cat the cat the cat the cat the cat the cat
greedy, penalty 1.5 : the cat sat on the mat
Nothing is broken. Greedy takes the best next token, the best token after
'cat' is 'the', and the best token after 'the' is 'cat'. The loop is the
correct answer to the question greedy decoding is asking.
what the penalty does to one step. after 'the cat', the scores from 'cat':
no penalty {'</s>': 0.0, 'sat': 3.0, 'the': 4.0}
penalty 1.5 {'</s>': 0.0, 'sat': 3.0, 'the': 2.667}
penalty 3.0 {'</s>': 0.0, 'sat': 3.0, 'the': 1.333}
The end marker at 0.0 does not move: the rule only touches tokens already
used. And 'the' is being punished for being a word English needs to repeat.
a penalty below 1 is refused: a penalty below 1 rewards repetition instead of discouraging it
ok greedy with no penalty runs into a two-token loop
ok and never reaches the end marker, it just hits the length limit
ok a penalty of 1.5 breaks the loop
ok and lets the sentence finish on its own
ok a penalty of 1.0 changes nothing at all
ok the penalty leaves tokens that have not been used alone
ok but it does punish a word that English needs to repeat
ok a penalty below one is refused, because it would reward repetition
Three variants exist and the difference is worth knowing.
repetition_penalty is the multiplicative rule above and is applied once,
however many times the token appeared. frequency_penalty subtracts an amount
proportional to the count, so the tenth repeat is punished ten times as hard.
presence_penalty subtracts a flat amount for having appeared at all. OpenAI
exposes the last two, Hugging Face the first. All three share the same defect the assertions
above make explicit: they punish tokens for being repeated, and they cannot tell a
degenerate loop from the word "the".
Because of that, a repetition penalty is the wrong tool for structured output. Turn one
on while generating JSON and you are penalising every repeated quote, brace and comma —
exactly the characters the format requires. If a model is looping, prefer sampling over
greedy, then no_repeat_ngram_size, which forbids repeating a whole n-gram
rather than a single token, and leave the penalty at 1.0 for anything that has a
syntax.
All four together, and the order they go in
One decode step is: logits, temperature, top-k, top-p, draw. That order is not decoration. Each filter acts on the distribution the previous step produced, so swapping two of them gives a different set of reachable tokens for the same settings.
# One decode step, all four settings, in the order a real library applies them.
import math
import random
TOKENS = ("cat", "dog", "bird", "fish", "quark", "ledger", "sonata", "kelp")
LOGITS = [3.0, 2.6, 2.2, 1.8, 1.4, 1.0, 0.6, 0.2]
def softmax(logits, temperature=1.0):
scaled = [z / temperature for z in logits]
shift = max(scaled)
exps = [math.exp(z - shift) for z in scaled]
total = sum(exps)
return [e / total for e in exps]
def top_k(probs, k):
ranked = sorted(range(len(probs)), key=lambda i: -probs[i])
keep = set(ranked[:k])
mass = sum(probs[i] for i in keep)
return [probs[i] / mass if i in keep else 0.0 for i in range(len(probs))]
def top_p(probs, p):
ranked = sorted(range(len(probs)), key=lambda i: -probs[i])
keep, running = [], 0.0
for i in ranked:
keep.append(i)
running += probs[i]
if running >= p:
break
mass = sum(probs[i] for i in keep)
kept = set(keep)
return [probs[i] / mass if i in kept else 0.0 for i in range(len(probs))]
def decode_step(logits, rng, temperature=1.0, k=None, p=None):
"""Temperature first, then top-k, then top-p, then draw."""
probs = softmax(logits, temperature)
if k is not None:
probs = top_k(probs, k)
if p is not None:
probs = top_p(probs, p)
r, running = rng.random(), 0.0
for i, x in enumerate(probs):
running += x
if r < running:
return i, probs
return len(probs) - 1, probs
def alive(dist):
return sum(1 for x in dist if x > 0.0)
def run(seed, n, **kw):
rng = random.Random(seed)
return [TOKENS[decode_step(LOGITS, rng, **kw)[0]] for _ in range(n)]
print(" settings reachable eight draws with seed 42")
for label, kw in (("nothing, plain sampling", {}),
("temperature 0.5", {"temperature": 0.5}),
("top_k 3", {"k": 3}),
("top_p 0.6", {"p": 0.6}),
("temperature 0.7 + top_k 5 + top_p 0.9",
{"temperature": 0.7, "k": 5, "p": 0.9})):
_, dist = decode_step(LOGITS, random.Random(0), **kw)
print(" {:37} {:9} {}".format(label, alive(dist), " ".join(run(42, 8, **kw))))
print()
print("order matters, and this is the trap. the same two settings, swapped:")
hot_then_cut = top_p(softmax(LOGITS, 2.0), 0.6)
cut_then_hot = softmax([z if x > 0.0 else -1e9
for z, x in zip(LOGITS, top_p(softmax(LOGITS), 0.6))], 2.0)
print(" temperature 2.0, then top_p 0.6 :", alive(hot_then_cut), "tokens reachable")
print(" top_p 0.6, then temperature 2.0 :", alive(cut_then_hot), "tokens reachable")
print("A high temperature flattens the distribution, so top-p then has to keep")
print("more tokens to reach the same mass. Filter first and it never can.")
_, plain = decode_step(LOGITS, random.Random(0))
_, both = decode_step(LOGITS, random.Random(0), temperature=0.7, k=5, p=0.9)
test("with no filter, every token in the vocabulary is reachable", alive(plain), 8)
test("the combined settings narrow it", alive(both) < 8, True)
test("and what is left is still a distribution", close(sum(both), 1.0), True)
test("the same seed and settings replay exactly", run(42, 8, k=3), run(42, 8, k=3))
test("a different seed gives a different sequence",
run(42, 8, k=3) == run(7, 8, k=3), False)
test("temperature before top-p leaves more tokens alive than after it",
alive(hot_then_cut) > alive(cut_then_hot), True)
test("greedy is a corner of this same function, at k = 1",
len(set(run(42, 8, k=1))), 1)What it printed, when the build ran it
settings reachable eight draws with seed 42 nothing, plain sampling 8 bird cat cat cat fish bird quark cat temperature 0.5 8 dog cat cat cat dog dog bird cat top_k 3 3 dog cat cat cat dog dog bird cat top_p 0.6 3 dog cat cat cat dog dog bird cat temperature 0.7 + top_k 5 + top_p 0.9 4 dog cat cat cat dog dog bird cat order matters, and this is the trap. the same two settings, swapped: temperature 2.0, then top_p 0.6 : 4 tokens reachable top_p 0.6, then temperature 2.0 : 3 tokens reachable A high temperature flattens the distribution, so top-p then has to keep more tokens to reach the same mass. Filter first and it never can. ok with no filter, every token in the vocabulary is reachable ok the combined settings narrow it ok and what is left is still a distribution ok the same seed and settings replay exactly ok a different seed gives a different sequence ok temperature before top-p leaves more tokens alive than after it ok greedy is a corner of this same function, at k = 1
- Factual work. Temperature 0, or greedy. You want the same answer twice, and you want the token the model actually believes.
- Anything with a syntax. Low temperature, no repetition penalty, and a validator afterwards. JSON is not a creative writing task.
- Open-ended prose. Temperature around 0.7 to 1.0 with
top_pnear 0.9, and akas a safety ceiling. - Reproducibility. Pin the seed and log the settings alongside the output. Without them a failing generation cannot be replayed, and "the model got worse" cannot be distinguished from "someone changed top_p".
What a real project would use instead
Every one of these functions exists in the library, applied over a whole vocabulary in compiled code, in the order above. The reason to have written them once is that the parameter names in the API are now transparent rather than magic, and when a generation goes strange you know which filter to suspect.
# The same four settings through a real generation API.
# NOT RUN HERE - torch and transformers are not installed.
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
tok = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")
ids = tok("the cat sat on", return_tensors="pt").input_ids
# Greedy. No sampling, no seed needed, same answer every time.
model.generate(ids, do_sample=False, max_new_tokens=20)
# Sampling with all four filters. do_sample=True is the switch people forget:
# leave it False and temperature, top_k and top_p are silently ignored.
torch.manual_seed(42)
model.generate(
ids,
do_sample=True,
temperature=0.8,
top_k=50, # hard ceiling
top_p=0.9, # adaptive width
repetition_penalty=1.1, # leave at 1.0 for JSON and code
no_repeat_ngram_size=3, # a better answer to loops than the penalty
max_new_tokens=60,
)
# The filters as standalone objects, applied in this order, which is the order
# the plain-Python decode_step above uses:
from transformers import (TemperatureLogitsWarper, TopKLogitsWarper,
TopPLogitsWarper, RepetitionPenaltyLogitsProcessor)
# And the same four settings over an HTTP API, where seed may be a hint rather
# than a guarantee, because batching on the server can change the arithmetic:
# {"temperature": 0.8, "top_p": 0.9, "frequency_penalty": 0.2, "seed": 42}torch and transformers are not installed in this environment, so this panel was never executed. Every filter it names was written out in plain Python above and asserted there; this shows only the parameter names a real API uses for them.
One thing to carry into an interview: none of this changes the model. Decoding is a pure function of the logits, so an unhappy output is either a bad distribution or a bad choice from a fine distribution, and the two have completely different fixes. Log the top few logits for a step that went wrong and you can tell which. If the right token was ranked second, that is a decoding problem and cheap to fix. If it was not in the top fifty, no setting on this page will help you.
20 Attention, from scratch
Lesson 16 ended with two failures: one vector per word cannot handle context, and counting cannot see word order. Attention fixes the first. Each token builds its own new vector by asking every other token a question, scoring the answers, and taking a weighted average of what they offer. Three steps, one softmax, no learned machinery beyond three matrices.
Query, key, value, and the three lines that use them
The vocabulary is worth ten seconds. A query is what a token is looking for. A key is what a token advertises about itself. A value is what a token hands over if you attend to it. Score every query against every key with a dot product, softmax each row, and mix the values in those proportions. Everything below is hand-written so the answers are checkable.
# Scaled dot-product attention, on four tokens and vectors small enough to read.
import math
TOKENS = ["the", "cat", "drank", "milk"]
AXES = ["animal", "action", "food", "filler"]
# Keys and values: what each token is. One axis lit per token, by hand.
K = [[0.0, 0.0, 0.0, 1.0], # the - a filler word
[1.0, 0.0, 0.0, 0.0], # cat - an animal
[0.0, 1.0, 0.0, 0.0], # drank - an action
[0.0, 0.0, 1.0, 0.0]] # milk - a food
V = [row[:] for row in K] # values are the same content, kept readable
# Queries: what each token is looking for. Also by hand, so the answers are checkable.
Q = [[0.0, 0.0, 0.0, 6.0], # the - looking for filler, i.e. itself
[0.0, 6.0, 0.0, 0.0], # cat - looking for the action
[6.0, 0.0, 6.0, 0.0], # drank - looking for who did it and to what
[0.0, 6.0, 0.0, 0.0]] # milk - looking for the action
D_K = len(AXES)
def dot(a, b):
return sum(x * y for x, y in zip(a, b))
def softmax(row):
shift = max(row) # lesson 6's subtract-the-max
exps = [math.exp(v - shift) for v in row]
total = sum(exps)
return [e / total for e in exps]
def attention(queries, keys, values):
"""Three lines of arithmetic. Everything else in this lesson is a variation."""
scale = math.sqrt(len(keys[0]))
scores = [[dot(q, k) / scale for k in keys] for q in queries]
weights = [softmax(row) for row in scores]
out = [[sum(w[j] * values[j][c] for j in range(len(values))) for c in range(len(values[0]))]
for w in weights]
return scores, weights, out
scores, weights, out = attention(Q, K, V)
print("axes:", AXES)
print()
print("raw scores, query row against key column, before scaling")
print(" " + "".join("{:>9}".format(t) for t in TOKENS))
for i, t in enumerate(TOKENS):
print(" {:8} ".format(t) + "".join("{:9.4f}".format(dot(Q[i], K[j]))
for j in range(len(TOKENS))))
print()
print("the same scores divided by sqrt(d_k) = sqrt({}) = {:.1f}".format(D_K, math.sqrt(D_K)))
print(" " + "".join("{:>9}".format(t) for t in TOKENS))
for i, t in enumerate(TOKENS):
print(" {:8} ".format(t) + "".join("{:9.4f}".format(v) for v in scores[i]))
print()
print("ATTENTION WEIGHTS after softmax on each row - who looked at whom")
print(" " + "".join("{:>9}".format(t) for t in TOKENS))
for i, t in enumerate(TOKENS):
print(" {:8} ".format(t) + "".join("{:9.4f}".format(v) for v in weights[i]))
print()
print("Read row by row. 'cat' put 87% of its attention on 'drank'. 'drank' split")
print("its attention almost evenly between 'cat' and 'milk' - the two things it")
print("was asking about. 'the' found nothing interesting and looked at itself.")
print()
print("the output, one row per token: the values, mixed in those proportions")
print(" " + "".join("{:>9}".format(a) for a in AXES))
for i, t in enumerate(TOKENS):
print(" {:8} ".format(t) + "".join("{:9.4f}".format(v) for v in out[i]))
rows_sum = [sum(row) for row in weights]
drank = TOKENS.index("drank")
print()
print("every row of the weight matrix sums to", ["{:.10f}".format(v) for v in rows_sum])
test("every row of the attention weights sums to one",
all(close(v, 1.0) for v in rows_sum), True)
test("no weight is negative and none exceeds one",
all(0.0 < v < 1.0 for row in weights for v in row), True)
test("cat put most of its attention on drank",
weights[TOKENS.index("cat")].index(max(weights[TOKENS.index("cat")])), drank)
test("drank split its attention evenly between cat and milk",
[close(weights[drank][1], weights[drank][3], 1e-12), weights[drank][1] > 0.45],
[True, True])
test("and gave almost nothing to the filler word or to itself",
[weights[drank][0] < 0.03, weights[drank][2] < 0.03], [True, True])
test("the output for cat is mostly the value row for drank",
[close(out[1][1], 0.87, 0.01), out[1][1] > 3.0 * out[1][0]], [True, True])
test("every output is a weighted average of the values, so it stays inside their range",
all(min(row[c] for row in V) <= out[i][c] <= max(row[c] for row in V)
for i in range(len(TOKENS)) for c in range(len(AXES))), True)What it printed, when the build ran it
axes: ['animal', 'action', 'food', 'filler']
raw scores, query row against key column, before scaling
the cat drank milk
the 6.0000 0.0000 0.0000 0.0000
cat 0.0000 0.0000 6.0000 0.0000
drank 0.0000 6.0000 0.0000 6.0000
milk 0.0000 0.0000 6.0000 0.0000
the same scores divided by sqrt(d_k) = sqrt(4) = 2.0
the cat drank milk
the 3.0000 0.0000 0.0000 0.0000
cat 0.0000 0.0000 3.0000 0.0000
drank 0.0000 3.0000 0.0000 3.0000
milk 0.0000 0.0000 3.0000 0.0000
ATTENTION WEIGHTS after softmax on each row - who looked at whom
the cat drank milk
the 0.8700 0.0433 0.0433 0.0433
cat 0.0433 0.0433 0.8700 0.0433
drank 0.0237 0.4763 0.0237 0.4763
milk 0.0433 0.0433 0.8700 0.0433
Read row by row. 'cat' put 87% of its attention on 'drank'. 'drank' split
its attention almost evenly between 'cat' and 'milk' - the two things it
was asking about. 'the' found nothing interesting and looked at itself.
the output, one row per token: the values, mixed in those proportions
animal action food filler
the 0.0433 0.0433 0.0433 0.8700
cat 0.0433 0.8700 0.0433 0.0433
drank 0.4763 0.0237 0.4763 0.0237
milk 0.0433 0.8700 0.0433 0.0433
every row of the weight matrix sums to ['1.0000000000', '1.0000000000', '1.0000000000', '1.0000000000']
ok every row of the attention weights sums to one
ok no weight is negative and none exceeds one
ok cat put most of its attention on drank
ok drank split its attention evenly between cat and milk
ok and gave almost nothing to the filler word or to itself
ok the output for cat is mostly the value row for drank
ok every output is a weighted average of the values, so it stays inside their range
Four things about that matrix are worth naming, because they are what people get asked:
- it is not symmetric. "cat" gives "drank" 0.87 while "drank" gives "cat" 0.48. Attention is a directed question, not a similarity measure.
- every row sums to 1, so each output is a weighted average of the values and can never leave the range the values span. That is asserted above.
- the output has one row per token, same shape as the input. Attention is a rewrite, not a reduction, which is why blocks stack.
- the cost is the square of the sequence length. Four tokens means sixteen scores; a thousand tokens means a million. That single fact drives the entire long-context research literature.
In a real model Q, K and V are not hand-written: each
is the token embeddings multiplied by a learned matrix, so the model decides for itself what
it is worth asking about. The panel above skips those three matrix multiplications so the
arithmetic in the middle is visible; the multi-head panel below puts them back.
Why divide by the square root of d_k
The scaling looks like a detail and is the difference between a block that trains and one
that does not. A dot product over d_k terms grows with d_k, and a
softmax fed large numbers stops being an average and becomes a lookup — with almost no
gradient left. Both halves of that are measured below rather than claimed.
# Why the / sqrt(d_k). Without it the softmax saturates and the gradient dies.
import math
import random
def dot(a, b):
return sum(x * y for x, y in zip(a, b))
def softmax(row):
shift = max(row)
exps = [math.exp(v - shift) for v in row]
total = sum(exps)
return [e / total for e in exps]
def entropy(p):
"""How spread out the attention is. Zero means it looked at exactly one token."""
return -sum(v * math.log(v) for v in p if v > 0.0)
DIMS = (4, 16, 64, 256)
KEYS = 12
TRIALS = 120
def experiment(d_k):
"""Draw a query and KEYS keys at random, TRIALS times, and average the results."""
rng = random.Random(42)
raw_max = raw_ent = scaled_max = scaled_ent = 0.0
raw_slope = scaled_slope = 0.0
collapsed = 0
spread_of_scores = []
same_winner = True
for _ in range(TRIALS):
q = [rng.gauss(0.0, 1.0) for _ in range(d_k)]
keys = [[rng.gauss(0.0, 1.0) for _ in range(d_k)] for _ in range(KEYS)]
raw = [dot(q, k) for k in keys]
spread_of_scores.extend(raw)
scaled = [v / math.sqrt(d_k) for v in raw]
p_raw = softmax(raw)
p_scaled = softmax(scaled)
raw_max += max(p_raw) / TRIALS
raw_ent += entropy(p_raw) / TRIALS
scaled_max += max(p_scaled) / TRIALS
scaled_ent += entropy(p_scaled) / TRIALS
# p(1 - p) is the slope of a softmax output against its own score: what
# backprop has to multiply by when it passes back through the attention.
raw_slope += max(p_raw) * (1.0 - max(p_raw)) / TRIALS
scaled_slope += max(p_scaled) * (1.0 - max(p_scaled)) / TRIALS
if max(p_raw) > 0.99:
collapsed += 1
if p_raw.index(max(p_raw)) != p_scaled.index(max(p_scaled)):
same_winner = False
mean = sum(spread_of_scores) / len(spread_of_scores)
sd = math.sqrt(sum((v - mean) ** 2 for v in spread_of_scores) / len(spread_of_scores))
return {"raw_max": raw_max, "raw_ent": raw_ent, "scaled_max": scaled_max,
"scaled_ent": scaled_ent, "sd": sd, "same_winner": same_winner,
"raw_slope": raw_slope, "scaled_slope": scaled_slope,
"collapsed": collapsed / float(TRIALS)}
print("Each component of q and k is a random number with a spread of about 1. The dot")
print("product adds up d_k of their products, and adding up d_k independent things")
print("multiplies the spread by sqrt(d_k). Measured over {} random queries against".format(TRIALS))
print("{} keys each:".format(KEYS))
print()
print(" d_k spread of q.k sqrt(d_k) spread after dividing")
results = {}
for d_k in DIMS:
r = experiment(d_k)
results[d_k] = r
print(" {:5} {:13.4f} {:9.4f} {:21.4f}".format(
d_k, r["sd"], math.sqrt(d_k), r["sd"] / math.sqrt(d_k)))
print()
print("Now what that does to the softmax over those {} keys, averaged over the".format(KEYS))
print("{} trials. 'biggest' is how much attention went to the favourite key, and".format(TRIALS))
print("entropy is how much was left over for everyone else. Uniform would be {:.4f}.".format(
math.log(KEYS)))
print()
print(" d_k unscaled biggest entropy collapsed scaled biggest entropy")
for d_k in DIMS:
r = results[d_k]
print(" {:5} {:17.6f} {:7.4f} {:8.0%} {:15.6f} {:7.4f}".format(
d_k, r["raw_max"], r["raw_ent"], r["collapsed"], r["scaled_max"], r["scaled_ent"]))
big = results[256]
print()
print("'collapsed' is the share of trials where one key took more than 99% of the")
print("attention. At d_k = 256 that is {:.0%} of them - the weighted average has become".format(
big["collapsed"]))
print("a lookup. Scaled, the favourite averages {:.4f} and the rest still contribute.".format(
big["scaled_max"]))
slope_raw = big["raw_slope"]
slope_scaled = big["scaled_slope"]
print()
print("mean slope of the unscaled softmax : p(1-p) = {:.3e}".format(slope_raw))
print("mean slope of the scaled one : p(1-p) = {:.3e}".format(slope_scaled))
print("So {:.0f} times more gradient reaches the layers below.".format(slope_scaled / slope_raw))
print()
print("The division is not cosmetic and it is not a normalisation. It keeps the scores")
print("in the range where a softmax still has a derivative - which is the same")
print("saturation trap lesson 15 hit with too large a learning rate.")
test("the spread of a raw dot product really does grow like sqrt(d_k)",
all(close(results[d]["sd"] / math.sqrt(d), 1.0, 0.1) for d in DIMS), True)
test("dividing by sqrt(d_k) brings every dimension back to about one",
close(results[256]["sd"] / 16.0, results[4]["sd"] / 2.0, 0.15), True)
test("unscaled, attention sharpens with every step up in d_k",
[results[4]["raw_max"] < results[64]["raw_max"],
results[64]["raw_max"] < results[256]["raw_max"]], [True, True])
test("and by d_k = 256 most trials have collapsed onto a single key",
[results[4]["collapsed"] < 0.05, results[256]["collapsed"] > 0.5], [True, True])
test("scaled, the sharpness barely moves as d_k grows",
close(results[256]["scaled_max"], results[4]["scaled_max"], 0.1), True)
test("so it passes back under a third of the gradient the scaled version does",
[slope_raw < 0.1, slope_scaled > 3.0 * slope_raw], [True, True])
test("and scaling never changes which key wins, only how sharply",
all(results[d]["same_winner"] for d in DIMS), True)What it printed, when the build ran it
Each component of q and k is a random number with a spread of about 1. The dot
product adds up d_k of their products, and adding up d_k independent things
multiplies the spread by sqrt(d_k). Measured over 120 random queries against
12 keys each:
d_k spread of q.k sqrt(d_k) spread after dividing
4 1.9527 2.0000 0.9763
16 3.9789 4.0000 0.9947
64 8.2096 8.0000 1.0262
256 16.1765 16.0000 1.0110
Now what that does to the softmax over those 12 keys, averaged over the
120 trials. 'biggest' is how much attention went to the favourite key, and
entropy is how much was left over for everyone else. Uniform would be 2.4849.
d_k unscaled biggest entropy collapsed scaled biggest entropy
4 0.452636 1.6080 0% 0.263469 2.1417
16 0.680288 0.8817 6% 0.271107 2.1227
64 0.889003 0.2887 39% 0.312375 2.0535
256 0.927592 0.1938 56% 0.281413 2.0894
'collapsed' is the share of trials where one key took more than 99% of the
attention. At d_k = 256 that is 56% of them - the weighted average has become
a lookup. Scaled, the favourite averages 0.2814 and the rest still contribute.
mean slope of the unscaled softmax : p(1-p) = 5.063e-02
mean slope of the scaled one : p(1-p) = 1.932e-01
So 4 times more gradient reaches the layers below.
The division is not cosmetic and it is not a normalisation. It keeps the scores
in the range where a softmax still has a derivative - which is the same
saturation trap lesson 15 hit with too large a learning rate.
ok the spread of a raw dot product really does grow like sqrt(d_k)
ok dividing by sqrt(d_k) brings every dimension back to about one
ok unscaled, attention sharpens with every step up in d_k
ok and by d_k = 256 most trials have collapsed onto a single key
ok scaled, the sharpness barely moves as d_k grows
ok so it passes back under a third of the gradient the scaled version does
ok and scaling never changes which key wins, only how sharply
The reasoning behind the constant is worth being able to reproduce on a whiteboard,
because it is a standard interview question. If the components of q and
k are independent with variance 1, then each product term has variance 1, and
d_k of them added together has variance d_k — so the score has a
spread of sqrt(d_k). Dividing by sqrt(d_k) restores a spread of
about 1, which is exactly what the third column above measures. Nothing about the number is
arbitrary and nothing about it is learned.
Causal masking
A language model predicts the next token. If position 2 can attend to position 3 then it
can read the answer, and the training loss will be excellent and the model useless. The fix
is one if: set the score to minus infinity wherever the key is in the query's
future.
# Causal masking: stop a token from reading the answer off the tokens after it.
import math
TOKENS = ["the", "cat", "drank", "milk"]
K = [[0.0, 0.0, 0.0, 1.0],
[1.0, 0.0, 0.0, 0.0],
[0.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 1.0, 0.0]]
V = [row[:] for row in K]
Q = [[0.0, 0.0, 0.0, 6.0],
[0.0, 6.0, 0.0, 0.0],
[6.0, 0.0, 6.0, 0.0],
[0.0, 6.0, 0.0, 0.0]]
NEG = float("-inf")
def dot(a, b):
return sum(x * y for x, y in zip(a, b))
def softmax(row):
shift = max(row) # always finite: the diagonal is never masked
exps = [0.0 if v == NEG else math.exp(v - shift) for v in row]
total = sum(exps)
return [e / total for e in exps]
def attention(queries, keys, values, causal=False):
scale = math.sqrt(len(keys[0]))
weights = []
for i, q in enumerate(queries):
row = []
for j, k in enumerate(keys):
if causal and j > i:
row.append(NEG) # position i may not see position j
else:
row.append(dot(q, k) / scale)
weights.append(softmax(row))
out = [[sum(w[j] * values[j][c] for j in range(len(values))) for c in range(len(values[0]))]
for w in weights]
return weights, out
open_weights, open_out = attention(Q, K, V)
causal_weights, causal_out = attention(Q, K, V, causal=True)
def show(title, matrix):
print(title)
print(" " + "".join("{:>9}".format(t) for t in TOKENS))
for i, t in enumerate(TOKENS):
print(" {:8} ".format(t) + "".join("{:9.4f}".format(v) for v in matrix[i]))
show("everyone sees everyone - what an encoder does", open_weights)
print()
show("causal: position i may only look at positions up to i", causal_weights)
print()
print("The upper triangle is exactly zero, not merely small, because the score there")
print("was minus infinity and exp(-inf) is 0. The rows still sum to 1, so each token")
print("is still taking a weighted average - just over a shorter list.")
print()
print("Row 'the' is [1, 0, 0, 0]: the first token has nothing to look at but itself.")
print("Row 'cat' has to settle for a 50-50 split with 'the', because the word it was")
print("actually asking for - 'drank' - is in its future. Row 'drank' gives 'cat' 0.909")
print("where the open version gave 0.476: the mass that went to 'milk' did not")
print("disappear, it was redistributed when the row was renormalised.")
# The property that makes this worth doing. Change the last token and the earlier
# outputs must not move, or training on every position at once would be cheating.
K2 = [row[:] for row in K]
V2 = [row[:] for row in V]
K2[3] = [0.0, 0.0, 0.0, 0.0] # replace 'milk' with something else entirely
V2[3] = [9.0, 9.0, 9.0, 9.0]
_, changed_out = attention(Q, K2, V2, causal=True)
_, changed_open = attention(Q, K2, V2)
print()
print("now replace the last token's key and value with something else:")
print(" causal, output row for 'cat' before :",
["{:.4f}".format(v) for v in causal_out[1]])
print(" causal, output row for 'cat' after :",
["{:.4f}".format(v) for v in changed_out[1]])
print(" open, output row for 'cat' after :",
["{:.4f}".format(v) for v in changed_open[1]])
print()
print("That is the whole reason for the mask. Because no earlier position can see a")
print("later one, a single forward pass over a sentence gives you a valid prediction")
print("at every position at once - four training examples for the price of one.")
upper = [causal_weights[i][j] for i in range(4) for j in range(4) if j > i]
rows_sum = [sum(row) for row in causal_weights]
earlier_unchanged = all(close(a, b) for i in (0, 1, 2)
for a, b in zip(causal_out[i], changed_out[i]))
test("every masked weight is exactly zero, not just small", set(upper), {0.0})
test("the rows of the masked matrix still each sum to one",
all(close(v, 1.0) for v in rows_sum), True)
test("the first token can only attend to itself", causal_weights[0], [1.0, 0.0, 0.0, 0.0])
test("the last token is the only one that still sees the whole sentence",
[sum(1 for v in row if v > 0.0) for row in causal_weights], [1, 2, 3, 4])
test("changing the last token leaves every earlier output untouched",
earlier_unchanged, True)
test("without the mask it changes all of them, which would be reading the answer",
all(close(a, b) for a, b in zip(open_out[1], changed_open[1])), False)What it printed, when the build ran it
everyone sees everyone - what an encoder does
the cat drank milk
the 0.8700 0.0433 0.0433 0.0433
cat 0.0433 0.0433 0.8700 0.0433
drank 0.0237 0.4763 0.0237 0.4763
milk 0.0433 0.0433 0.8700 0.0433
causal: position i may only look at positions up to i
the cat drank milk
the 1.0000 0.0000 0.0000 0.0000
cat 0.5000 0.5000 0.0000 0.0000
drank 0.0453 0.9094 0.0453 0.0000
milk 0.0433 0.0433 0.8700 0.0433
The upper triangle is exactly zero, not merely small, because the score there
was minus infinity and exp(-inf) is 0. The rows still sum to 1, so each token
is still taking a weighted average - just over a shorter list.
Row 'the' is [1, 0, 0, 0]: the first token has nothing to look at but itself.
Row 'cat' has to settle for a 50-50 split with 'the', because the word it was
actually asking for - 'drank' - is in its future. Row 'drank' gives 'cat' 0.909
where the open version gave 0.476: the mass that went to 'milk' did not
disappear, it was redistributed when the row was renormalised.
now replace the last token's key and value with something else:
causal, output row for 'cat' before : ['0.5000', '0.0000', '0.0000', '0.5000']
causal, output row for 'cat' after : ['0.5000', '0.0000', '0.0000', '0.5000']
open, output row for 'cat' after : ['0.4332', '1.2599', '0.3899', '0.4332']
That is the whole reason for the mask. Because no earlier position can see a
later one, a single forward pass over a sentence gives you a valid prediction
at every position at once - four training examples for the price of one.
ok every masked weight is exactly zero, not just small
ok the rows of the masked matrix still each sum to one
ok the first token can only attend to itself
ok the last token is the only one that still sees the whole sentence
ok changing the last token leaves every earlier output untouched
ok without the mask it changes all of them, which would be reading the answer
Two mask bugs, both of which produce a model that looks fine. Using a large negative
number instead of -inf — say -1e9 — mostly works and then
overflows to nan in float16, which is why real implementations use the dtype's
minimum rather than a round number. And masking with >= instead of
> stops a token attending to itself, which quietly removes the diagonal and
makes the first row a softmax over nothing. Neither raises. The way to catch both is the
assertion in the panel: print the matrix and check the rows sum to 1.
Multi-head attention
One attention produces one opinion about which tokens matter. A sentence usually has several relationships in it at once — who did what, which noun a pronoun refers to, where the clause ends — and averaging them into one weight matrix loses all of them. So run the whole thing several times on narrower slices and glue the results back together.
# Multi-head attention: several attentions side by side, then glued back together.
import math
import random
TOKENS = ["the", "cat", "drank", "milk"]
X = [[0.0, 0.0, 0.0, 1.0], # the token embeddings from the first panel
[1.0, 0.0, 0.0, 0.0],
[0.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 1.0, 0.0]]
D_MODEL = 4
HEADS = 2
D_HEAD = D_MODEL // HEADS # 2. Each head works in half the width.
def dot(a, b):
return sum(x * y for x, y in zip(a, b))
def softmax(row):
shift = max(row)
exps = [math.exp(v - shift) for v in row]
total = sum(exps)
return [e / total for e in exps]
def matmul(rows, matrix):
"""rows is n by p, matrix is p by q, the answer is n by q."""
q = len(matrix[0])
return [[sum(r[k] * matrix[k][c] for k in range(len(matrix))) for c in range(q)]
for r in rows]
def projection(n_in, n_out, rng):
"""A learned matrix. Seeded, so this page prints the same numbers every build."""
return [[rng.uniform(-1.5, 1.5) for _ in range(n_out)] for _ in range(n_in)]
def attention(queries, keys, values):
scale = math.sqrt(len(keys[0]))
weights = [softmax([dot(q, k) / scale for k in keys]) for q in queries]
out = [[sum(w[j] * values[j][c] for j in range(len(values))) for c in range(len(values[0]))]
for w in weights]
return weights, out
rng = random.Random(42)
HEAD_PARAMS = []
for _ in range(HEADS):
HEAD_PARAMS.append((projection(D_MODEL, D_HEAD, rng), # W_q for this head
projection(D_MODEL, D_HEAD, rng), # W_k
projection(D_MODEL, D_HEAD, rng))) # W_v
W_O = projection(D_MODEL, D_MODEL, rng) # glues the heads back together
per_head_weights = []
per_head_out = []
for index, (w_q, w_k, w_v) in enumerate(HEAD_PARAMS):
q = matmul(X, w_q)
k = matmul(X, w_k)
v = matmul(X, w_v)
weights, out = attention(q, k, v)
per_head_weights.append(weights)
per_head_out.append(out)
print("head {}: its own Q, K and V, each {} wide".format(index, D_HEAD))
print(" " + "".join("{:>9}".format(t) for t in TOKENS))
for i, t in enumerate(TOKENS):
print(" {:8} ".format(t) + "".join("{:9.4f}".format(x) for x in weights[i]))
print()
# Concatenate: head 0's output followed by head 1's, back to D_MODEL wide.
concatenated = [per_head_out[0][i] + per_head_out[1][i] for i in range(len(TOKENS))]
output = matmul(concatenated, W_O)
print("concatenated, {} + {} = {} numbers per token".format(D_HEAD, D_HEAD, D_MODEL))
for i, t in enumerate(TOKENS):
print(" {:8} ".format(t) + "".join("{: 9.4f}".format(v) for v in concatenated[i]))
print()
print("after the output projection W_O, still {} wide, ready for the next layer".format(D_MODEL))
for i, t in enumerate(TOKENS):
print(" {:8} ".format(t) + "".join("{: 9.4f}".format(v) for v in output[i]))
# The two heads disagree about who matters, which is the entire point of having two.
gaps = [abs(per_head_weights[0][i][j] - per_head_weights[1][i][j])
for i in range(len(TOKENS)) for j in range(len(TOKENS))]
favourites = [[row.index(max(row)) for row in per_head_weights[h]] for h in range(HEADS)]
print()
print("head 0's favourite per row:", [TOKENS[i] for i in favourites[0]])
print("head 1's favourite per row:", [TOKENS[i] for i in favourites[1]])
print("biggest disagreement between the two matrices: {:.4f}".format(max(gaps)))
print()
# Cost. h heads of width d_model/h is the same parameter count as one head of d_model.
one_head = 3 * D_MODEL * D_MODEL
many_heads = HEADS * 3 * D_MODEL * D_HEAD
print("one head at full width : 3 * {} * {} = {} parameters".format(
D_MODEL, D_MODEL, one_head))
print("{} heads at width {} : {} * 3 * {} * {} = {} parameters".format(
HEADS, D_HEAD, HEADS, D_MODEL, D_HEAD, many_heads))
print("Identical. Splitting the width is free, and it buys several different views")
print("of the same sentence instead of one averaged compromise.")
test("each head produces a proper attention matrix, rows summing to one",
all(close(sum(row), 1.0) for h in range(HEADS) for row in per_head_weights[h]), True)
test("each head returns half the width", [len(per_head_out[h][0]) for h in range(HEADS)], [2, 2])
test("concatenating them restores the model width",
[len(concatenated[0]), len(output[0])], [D_MODEL, D_MODEL])
test("the two heads really do attend differently, or the second one is wasted",
max(gaps) > 0.2, True)
test("and they disagree about which token matters on at least one row",
favourites[0] == favourites[1], False)
test("splitting the width across heads costs exactly what one full-width head costs",
[one_head, many_heads], [48, 48])What it printed, when the build ran it
head 0: its own Q, K and V, each 2 wide
the cat drank milk
the 0.2442 0.5355 0.0924 0.1280
cat 0.1373 0.5323 0.1056 0.2248
drank 0.0916 0.3147 0.1822 0.4115
milk 0.4624 0.1838 0.2311 0.1226
head 1: its own Q, K and V, each 2 wide
the cat drank milk
the 0.3016 0.2847 0.1130 0.3007
cat 0.4200 0.1080 0.1057 0.3663
drank 0.1526 0.1285 0.5697 0.1493
milk 0.3417 0.2160 0.1154 0.3269
concatenated, 2 + 2 = 4 numbers per token
the -0.3629 -0.1696 -0.6576 -0.9013
cat -0.2080 -0.0220 -0.9846 -0.7873
drank 0.2388 -0.0355 -0.7570 -0.7563
milk -0.0487 -0.6980 -0.7827 -0.8569
after the output projection W_O, still 4 wide, ready for the next layer
the -0.7039 -0.9056 0.5760 -0.0006
cat -0.2541 -0.9520 0.9309 0.0940
drank -0.2417 -0.9814 0.5264 -0.3649
milk -0.0181 -1.7925 0.3530 -0.3776
head 0's favourite per row: ['cat', 'cat', 'milk', 'the']
head 1's favourite per row: ['the', 'the', 'drank', 'the']
biggest disagreement between the two matrices: 0.4243
one head at full width : 3 * 4 * 4 = 48 parameters
2 heads at width 2 : 2 * 3 * 4 * 2 = 48 parameters
Identical. Splitting the width is free, and it buys several different views
of the same sentence instead of one averaged compromise.
ok each head produces a proper attention matrix, rows summing to one
ok each head returns half the width
ok concatenating them restores the model width
ok the two heads really do attend differently, or the second one is wasted
ok and they disagree about which token matters on at least one row
ok splitting the width across heads costs exactly what one full-width head costs
The heads in that panel have random projections, so their disagreement is arbitrary. In a trained model it is not: heads specialise, and the specialisations are legible enough that people publish pictures of them — a head that tracks the previous token, a head that connects a verb to its subject, a head that finds the matching bracket. Nothing in the code makes that happen. It is what minimising next-token loss turns out to reward.
What a real implementation calls this
All four panels above are one function call in a framework, and the interesting part of the API is which of the arguments correspond to which panel.
# The same attention, as a real project writes it.
# NOT RUN HERE - torch is not installed in this environment.
import torch
import torch.nn as nn
import torch.nn.functional as F
# Panels one to three, in one call. The scaling and the mask are both arguments.
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
# Written out, it is the first panel with the batch dimensions kept:
scores = (q @ k.transpose(-2, -1)) / math.sqrt(q.size(-1)) # the / sqrt(d_k)
mask = torch.triu(torch.ones(n, n, dtype=torch.bool), diagonal=1)
scores = scores.masked_fill(mask, float("-inf")) # the causal mask
weights = scores.softmax(dim=-1) # rows summing to 1
out = weights @ v # the weighted average
# Panel four, packaged. embed_dim // num_heads is D_HEAD, and it must divide exactly.
attn = nn.MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True)
out, weights = attn(x, x, x, need_weights=True) # x three times is SELF attention
print(weights.shape) # (batch, n, n) - the matrix to plot
# Two things worth knowing that the panels above make obvious:
#
# need_weights=False is the default in the fused path, and the fast kernels
# (FlashAttention) never build the n by n matrix at all - which is why you cannot
# always get the picture out of a production model.
#
# The n by n matrix is why context length is expensive. 4k tokens is 16 million
# scores per head per layer; 32k tokens is 64 times that.torch is not installed in this environment, so this panel was never executed and q, k, v, x and n are never defined in it. Every line of arithmetic in it was written out in plain Python above and asserted: the scaling, the mask, the row sums, the weighted average and the per-head split all ran.
What to carry into the next lesson. Attention is: score with a dot product, divide by
sqrt(d_k), softmax each row, average the values. The output has the same shape
as the input, so it stacks. The mask is one if and it is what lets a model
train on every position of a sentence in a single pass. Multi-head is the same function on
narrower slices, at identical cost. And the one weakness left is the one lesson 16 ended on:
nothing in any of these panels knows what order the tokens came in — permute the input and
the output permutes with it. That is the first thing lesson 21 fixes.
21 A transformer block end to end
A transformer is one block repeated. The block is four things in a fixed order: attention, a residual add, a normalisation, and a small feed-forward network applied to each position on its own. Lesson 20 built the attention. This lesson builds everything around it and shows what each part is actually preventing.
Layer normalisation, and the two invariances it buys
Take one token's vector, subtract its mean, divide by its spread. That is the whole function. What it gains you is that the layer above stops caring how the layer below chose to scale things — which is worth demonstrating rather than describing.
# Layer normalisation, written out. Four lines and one guard.
import math
def layer_norm(x, gamma=None, beta=None, eps=1e-5):
"""Rescale one vector to mean 0 and spread 1, then let the model scale it back."""
n = float(len(x))
mean = sum(x) / n
variance = sum((v - mean) ** 2 for v in x) / n
denominator = math.sqrt(variance + eps) # eps, or a flat vector divides by zero
normalised = [(v - mean) / denominator for v in x]
if gamma is None:
return normalised
return [g * v + b for g, v, b in zip(gamma, normalised, beta)]
def mean_of(v):
return sum(v) / float(len(v))
def sd_of(v):
m = mean_of(v)
return math.sqrt(sum((x - m) ** 2 for x in v) / float(len(v)))
X = [2.0, -1.0, 4.0, 7.0]
normed = layer_norm(X)
print("input ", ["{: .4f}".format(v) for v in X])
print(" mean {:.4f}, spread {:.4f}".format(mean_of(X), sd_of(X)))
print()
print("layer_norm(x) ", ["{: .4f}".format(v) for v in normed])
print(" mean {:.10f}, spread {:.10f}".format(mean_of(normed), sd_of(normed)))
print()
print("Note what it normalised ACROSS: the four numbers of this one token. Not across")
print("the batch, not across the sequence. Every token is normalised on its own, which")
print("is why the answer does not change when the batch or the sentence length does.")
# The two invariances. Shift the input, or stretch it, and nothing downstream notices.
shifted = layer_norm([v + 100.0 for v in X])
stretched = layer_norm([v * 10.0 for v in X])
print()
print("layer_norm(x + 100)", ["{: .4f}".format(v) for v in shifted])
print("layer_norm(x * 10) ", ["{: .4f}".format(v) for v in stretched])
print("Both identical to the first. A layer below can drift in scale or offset as far")
print("as it likes; the layer above sees the same thing.")
# gamma and beta: the model's chance to undo the normalisation where it does not help.
GAMMA = [2.0, 2.0, 2.0, 2.0]
BETA = [1.0, 1.0, 1.0, 1.0]
affine = layer_norm(X, GAMMA, BETA)
print()
print("with gamma = 2, beta = 1:", ["{: .4f}".format(v) for v in affine])
print(" mean {:.4f}, spread {:.4f} <- shifted to beta, stretched to gamma".format(
mean_of(affine), sd_of(affine)))
# The guard. A vector with no spread at all is not a strange edge case: it is a padded
# row, a dead layer, or a feature that turned out constant.
FLAT = [3.0, 3.0, 3.0, 3.0]
print()
print("layer_norm on a flat vector:", ["{: .4f}".format(v) for v in layer_norm(FLAT)])
try:
layer_norm(FLAT, eps=0.0)
blew_up = False
except ZeroDivisionError as error:
blew_up = True
print("with eps = 0 it raises : ZeroDivisionError:", error)
test("the normalised vector has mean zero", close(mean_of(normed), 0.0), True)
test("and a spread of one", close(sd_of(normed), 1.0, 1e-5), True)
test("adding a hundred to every component changes nothing",
all(close(a, b) for a, b in zip(normed, shifted)), True)
test("multiplying every component by ten changes nothing either",
all(close(a, b, 1e-5) for a, b in zip(normed, stretched)), True)
test("gamma and beta put the mean and spread wherever the model wants them",
[close(mean_of(affine), 1.0), close(sd_of(affine), 2.0, 1e-5)], [True, True])
test("a flat vector normalises to zeros rather than to nan",
all(close(v, 0.0) for v in layer_norm(FLAT)), True)
test("and without eps that same vector divides by zero", blew_up, True)What it printed, when the build ran it
input [' 2.0000', '-1.0000', ' 4.0000', ' 7.0000'] mean 3.0000, spread 2.9155 layer_norm(x) ['-0.3430', '-1.3720', ' 0.3430', ' 1.3720'] mean 0.0000000000, spread 0.9999994118 Note what it normalised ACROSS: the four numbers of this one token. Not across the batch, not across the sequence. Every token is normalised on its own, which is why the answer does not change when the batch or the sentence length does. layer_norm(x + 100) ['-0.3430', '-1.3720', ' 0.3430', ' 1.3720'] layer_norm(x * 10) ['-0.3430', '-1.3720', ' 0.3430', ' 1.3720'] Both identical to the first. A layer below can drift in scale or offset as far as it likes; the layer above sees the same thing. with gamma = 2, beta = 1: [' 0.3140', '-1.7440', ' 1.6860', ' 3.7440'] mean 1.0000, spread 2.0000 <- shifted to beta, stretched to gamma layer_norm on a flat vector: [' 0.0000', ' 0.0000', ' 0.0000', ' 0.0000'] with eps = 0 it raises : ZeroDivisionError: float division by zero ok the normalised vector has mean zero ok and a spread of one ok adding a hundred to every component changes nothing ok multiplying every component by ten changes nothing either ok gamma and beta put the mean and spread wherever the model wants them ok a flat vector normalises to zeros rather than to nan ok and without eps that same vector divides by zero
Layer norm normalises across the features of one token. Batch norm, which came first, normalises across the batch instead — and that is why transformers do not use it. A batch statistic depends on which other examples happen to be in the batch, needs a running average kept for inference, and falls apart when sequences have different lengths and half the batch is padding. Layer norm has no batch dependence at all, which means one token gives the same answer whether it arrives alone or in a batch of 512. That property is worth more than the marginal difference in convergence.
What the residual connection is actually for
x + sublayer(x) looks like almost nothing. Its effect is on the backward pass:
the derivative of x + f(x) is 1 + f'(x), and that 1 is a
route from the loss back to the input that nothing can attenuate. Here it is measured on a
deep chain, with the analytic products checked against a nudge as lesson 14 insisted.
# What the residual connection is for. A deep chain, with and without it.
import math
W = 0.5
X0 = 0.5
def sublayer(x):
"""Stand-in for one attention or feed-forward block: something that bends."""
return W * math.tanh(x)
def d_sublayer(x):
"""Its local derivative, from lesson 13's tanh identity."""
return W * (1.0 - math.tanh(x) ** 2)
def run(depth, residual):
"""Forward through depth copies, keeping the derivative product as we go."""
x = X0
slope = 1.0
for _ in range(depth):
if residual:
slope *= 1.0 + d_sublayer(x) # d(x + f(x))/dx = 1 + f'(x)
x = x + sublayer(x)
else:
slope *= d_sublayer(x) # d f(x)/dx = f'(x)
x = sublayer(x)
return x, slope
print("one sublayer is f(x) = {} * tanh(x), starting from x = {}".format(W, X0))
print()
print(" depth plain: output d out / d in residual: output d out / d in")
plain = {}
residual = {}
for depth in (1, 2, 4, 8, 16, 32):
p_out, p_slope = run(depth, False)
r_out, r_slope = run(depth, True)
plain[depth] = p_slope
residual[depth] = r_slope
print(" {:5} {:17.10f} {:15.3e} {:17.4f} {:14.3e}".format(
depth, p_out, p_slope, r_out, r_slope))
print()
print("Plain, every layer multiplies the gradient by f'(x), which is less than 1, so")
print("after 32 layers there is nothing left for the first layer to learn from. This is")
print("the vanishing gradient, and it is why nobody trained a 30-layer network in 1995.")
print()
print("With a residual, every layer multiplies by 1 + f'(x). The 1 is the point: it is")
print("a route from the loss back to the input that no activation can attenuate.")
# The clearest version of the argument. Saturate the activation so f'(x) is zero.
def saturated(residual_path):
x = 30.0 # tanh(30) is 1.0 to every digit a float has
slope = 1.0
for _ in range(10):
local = W * (1.0 - math.tanh(x) ** 2)
slope *= (1.0 + local) if residual_path else local
x = (x + W * math.tanh(x)) if residual_path else W * math.tanh(x)
return slope
print()
print("now force f'(x) to zero by saturating tanh, ten layers deep:")
print(" plain d out / d in =", saturated(False), " <- the layer cannot learn at all")
print(" residual d out / d in =", saturated(True), " <- the gradient arrives intact")
# Check the analytic products against a nudge, as lesson 14 insisted.
def output_only(x, depth, residual_path):
for _ in range(depth):
x = (x + W * math.tanh(x)) if residual_path else W * math.tanh(x)
return x
H = 1e-6
checks = []
for residual_path in (False, True):
numeric = (output_only(X0 + H, 4, residual_path)
- output_only(X0 - H, 4, residual_path)) / (2.0 * H)
analytic = run(4, residual_path)[1]
checks.append(abs(numeric - analytic))
print()
print("depth 4, residual = {}: analytic {:.10f} nudge {:.10f}".format(
residual_path, analytic, numeric))
print()
print("And notice the residual column climbing. 1 + f'(x) is bigger than 1, so 32 of")
print("them multiply up rather than down. Residuals cure the vanishing gradient and")
print("introduce a growth problem, which is exactly the job the layer norm does.")
test("plain, the gradient shrinks with every extra layer",
[plain[2] < plain[1], plain[8] < plain[4], plain[32] < plain[16]], [True, True, True])
test("by 32 layers there is essentially none of it left", plain[32] < 1e-8, True)
test("with a residual the gradient at 32 layers is still large",
residual[32] > 1.0, True)
test("saturate the activation and plain delivers exactly zero, residual exactly one",
[saturated(False), saturated(True)], [0.0, 1.0])
test("both analytic products agree with a nudge on the whole chain",
all(gap < 1e-5 for gap in checks), True)
test("but the residual product grows with depth, which is what a norm is there to stop",
residual[32] > residual[1], True)What it printed, when the build ran it
one sublayer is f(x) = 0.5 * tanh(x), starting from x = 0.5
depth plain: output d out / d in residual: output d out / d in
1 0.2310585786 3.932e-01 0.7311 1.393e+00
2 0.1135163044 1.865e-01 1.0429 1.819e+00
4 0.0282277571 4.588e-02 1.8785 2.398e+00
8 0.0017636128 2.864e-03 3.8422 2.573e+00
16 0.0000068891 1.119e-05 7.8415 2.577e+00
32 0.0000000001 1.707e-10 15.8415 2.577e+00
Plain, every layer multiplies the gradient by f'(x), which is less than 1, so
after 32 layers there is nothing left for the first layer to learn from. This is
the vanishing gradient, and it is why nobody trained a 30-layer network in 1995.
With a residual, every layer multiplies by 1 + f'(x). The 1 is the point: it is
a route from the loss back to the input that no activation can attenuate.
now force f'(x) to zero by saturating tanh, ten layers deep:
plain d out / d in = 0.0 <- the layer cannot learn at all
residual d out / d in = 1.0 <- the gradient arrives intact
depth 4, residual = False: analytic 0.0458771505 nudge 0.0458771505
depth 4, residual = True: analytic 2.3983702611 nudge 2.3983702611
And notice the residual column climbing. 1 + f'(x) is bigger than 1, so 32 of
them multiply up rather than down. Residuals cure the vanishing gradient and
introduce a growth problem, which is exactly the job the layer norm does.
ok plain, the gradient shrinks with every extra layer
ok by 32 layers there is essentially none of it left
ok with a residual the gradient at 32 layers is still large
ok saturate the activation and plain delivers exactly zero, residual exactly one
ok both analytic products agree with a nudge on the whole chain
ok but the residual product grows with depth, which is what a norm is there to stop
The mental model worth keeping is the residual stream. The vector flowing down the middle of the model is not recomputed by each block; it is added to. Each block reads the stream, works out a correction, and writes that correction back. A block that has nothing useful to contribute can output near zero and the stream passes through unharmed, which is also why you can prune or skip blocks in a trained transformer and often lose surprisingly little.
The block, one step at a time
Now assemble it. Attention, residual, normalisation, feed-forward, residual. The panel prints every intermediate matrix so the order is visible, and then does the same block the other way round — normalise first, or normalise last — because that choice has consequences.
# One transformer block, assembled. The order of the four parts is the lesson.
import math
import random
TOKENS = ["the", "cat", "drank", "milk"]
D_MODEL = 4
D_FF = 16 # the feed-forward middle is conventionally 4 * d_model
NEG = float("-inf")
def layer_norm(x, gamma, beta, eps=1e-5):
n = float(len(x))
mean = sum(x) / n
variance = sum((v - mean) ** 2 for v in x) / n
denominator = math.sqrt(variance + eps)
return [g * ((v - mean) / denominator) + b for g, v, b in zip(gamma, x, beta)]
def norm_rows(rows, gamma, beta):
return [layer_norm(r, gamma, beta) for r in rows]
def matmul(rows, m):
return [[sum(r[k] * m[k][c] for k in range(len(m))) for c in range(len(m[0]))] for r in rows]
def softmax(row):
shift = max(row)
exps = [0.0 if v == NEG else math.exp(v - shift) for v in row]
total = sum(exps)
return [e / total for e in exps]
def attention(rows, w_q, w_k, w_v, w_o):
"""Lesson 20's causal self-attention, one head, with its output projection."""
q = matmul(rows, w_q)
k = matmul(rows, w_k)
v = matmul(rows, w_v)
scale = math.sqrt(len(k[0]))
out = []
for i in range(len(rows)):
scores = [NEG if j > i else sum(a * b for a, b in zip(q[i], k[j])) / scale
for j in range(len(rows))]
w = softmax(scores)
out.append([sum(w[j] * v[j][c] for j in range(len(rows))) for c in range(len(v[0]))])
return matmul(out, w_o)
def feed_forward(rows, w1, b1, w2, b2):
"""Same width in, same width out, wider in the middle. Position by position."""
out = []
for r in rows:
hidden = [max(0.0, sum(r[k] * w1[k][c] for k in range(len(w1))) + b1[c])
for c in range(len(w1[0]))]
out.append([sum(hidden[k] * w2[k][c] for k in range(len(w2))) + b2[c]
for c in range(len(w2[0]))])
return out
def matrix(n_in, n_out, rng):
return [[rng.uniform(-0.5, 0.5) for _ in range(n_out)] for _ in range(n_in)]
rng = random.Random(42)
W_Q, W_K, W_V, W_O = (matrix(D_MODEL, D_MODEL, rng) for _ in range(4))
W1 = matrix(D_MODEL, D_FF, rng)
B1 = [rng.uniform(-0.5, 0.5) for _ in range(D_FF)]
W2 = matrix(D_FF, D_MODEL, rng)
B2 = [rng.uniform(-0.5, 0.5) for _ in range(D_MODEL)]
G1 = [1.0] * D_MODEL # layer norm starts as the identity: gamma 1, beta 0
BE1 = [0.0] * D_MODEL
G2 = [1.0] * D_MODEL
BE2 = [0.0] * D_MODEL
X = [[0.2, -1.4, 0.9, 3.1],
[1.0, 0.1, -0.3, 0.0],
[-2.0, 0.5, 0.4, 1.2],
[0.3, 0.3, -1.1, 0.7]]
def add(a, b, scale=1.0):
return [[p + scale * q for p, q in zip(ra, rb)] for ra, rb in zip(a, b)]
def pre_norm_block(x, scale=1.0):
"""normalise, sublayer, add. What GPT-2 onwards actually does."""
h = add(x, attention(norm_rows(x, G1, BE1), W_Q, W_K, W_V, W_O), scale)
return add(h, feed_forward(norm_rows(h, G2, BE2), W1, B1, W2, B2), scale)
def post_norm_block(x, scale=1.0):
"""sublayer, add, normalise. What the original 2017 paper published."""
h = norm_rows(add(x, attention(x, W_Q, W_K, W_V, W_O), scale), G1, BE1)
return norm_rows(add(h, feed_forward(h, W1, B1, W2, B2), scale), G2, BE2)
def show(title, rows):
print(title)
for name, row in zip(TOKENS, rows):
print(" {:8} ".format(name) + "".join("{: 10.4f}".format(v) for v in row))
show("input, one row per token, {} numbers each".format(D_MODEL), X)
print()
# The four stages of a pre-norm block, in order, so nothing is hidden.
step1 = norm_rows(X, G1, BE1)
step2 = attention(step1, W_Q, W_K, W_V, W_O)
step3 = add(X, step2)
step4 = norm_rows(step3, G2, BE2)
step5 = feed_forward(step4, W1, B1, W2, B2)
step6 = add(step3, step5)
show("1. layer_norm(x)", step1)
print()
show("2. attention of that - tokens mixing, the only place they ever do", step2)
print()
show("3. x + attention, the first residual", step3)
print()
show("4. layer_norm again", step4)
print()
show("5. feed_forward of that - each token on its own, no mixing", step5)
print()
show("6. + the second residual: the block output", step6)
print()
print("Only step 2 lets one token see another. Steps 4 to 6 treat every row separately,")
print("which is why the feed-forward is described as position-wise and why two thirds of")
print("a transformer's parameters do no mixing at all.")
post = post_norm_block(X)
print()
show("the same block written post-norm: ln(x + sublayer(x))", post)
def rms(rows):
flat = [v for r in rows for v in r]
return math.sqrt(sum(v * v for v in flat) / len(flat))
print()
print("Stacking the same block over and over (weights reused, for a clean comparison):")
print()
print(" depth pre-norm stream size post-norm stream size")
pre_sizes = {}
post_sizes = {}
pre_state = [r[:] for r in X]
post_state = [r[:] for r in X]
for depth in range(13):
if depth in (0, 1, 2, 4, 8, 12):
pre_sizes[depth] = rms(pre_state)
post_sizes[depth] = rms(post_state)
print(" {:5} {:20.4f} {:21.4f}".format(depth, pre_sizes[depth], post_sizes[depth]))
pre_state = pre_norm_block(pre_state)
post_state = post_norm_block(post_state)
silent = pre_norm_block(X, scale=0.0)
params = 4 * D_MODEL * D_MODEL + D_MODEL * D_FF + D_FF + D_FF * D_MODEL + D_MODEL + 4 * D_MODEL
print()
print("pre-norm leaves the residual stream free to grow, so every large model ends with")
print("one final layer_norm before the output head. post-norm normalises last, so the")
print("stream is pinned - and the gradient has to pass through a normalisation on its")
print("way back to every layer, which is why deep post-norm stacks need a warmup")
print("schedule to train at all.")
print()
print("parameters in this block:", params)
test("the block returns exactly the shape it was given",
[len(step6), len(step6[0])], [len(TOKENS), D_MODEL])
test("silence both sublayers and the pre-norm block hands back its input untouched",
all(close(a, b) for ra, rb in zip(silent, X) for a, b in zip(ra, rb)), True)
test("post-norm ends on a normalisation, so every output row has spread one",
all(close(math.sqrt(sum((v - sum(r) / D_MODEL) ** 2 for v in r) / D_MODEL), 1.0, 1e-4)
for r in post), True)
test("pre-norm does not, which is why a final norm sits before the output head",
all(close(math.sqrt(sum((v - sum(r) / D_MODEL) ** 2 for v in r) / D_MODEL), 1.0, 1e-4)
for r in step6), False)
test("stacked twelve deep, the pre-norm stream has grown a long way",
pre_sizes[12] > 3.0 * pre_sizes[0], True)
test("while the post-norm stream is pinned wherever the normalisation puts it",
close(post_sizes[12], post_sizes[1], 0.2), True)
test("the block holds 228 parameters, 148 of them in the feed-forward",
[params, D_MODEL * D_FF + D_FF + D_FF * D_MODEL + D_MODEL], [228, 148])What it printed, when the build ran it
input, one row per token, 4 numbers each
the 0.2000 -1.4000 0.9000 3.1000
cat 1.0000 0.1000 -0.3000 0.0000
drank -2.0000 0.5000 0.4000 1.2000
milk 0.3000 0.3000 -1.1000 0.7000
1. layer_norm(x)
the -0.3092 -1.2986 0.1237 1.4841
cat 1.6502 -0.2063 -1.0314 -0.4126
drank -1.6748 0.3929 0.3102 0.9718
milk 0.3656 0.3656 -1.6819 0.9506
2. attention of that - tokens mixing, the only place they ever do
the 0.1515 -0.4971 0.3524 0.2595
cat -0.0066 0.0518 0.2664 0.2357
drank -0.0201 0.0595 0.1738 0.1666
milk -0.0811 0.1856 0.1572 0.1943
3. x + attention, the first residual
the 0.3515 -1.8971 1.2524 3.3595
cat 0.9934 0.1518 -0.0336 0.2357
drank -2.0201 0.5595 0.5738 1.3666
milk 0.2189 0.4856 -0.9428 0.8943
4. layer_norm again
the -0.2201 -1.4124 0.2576 1.3749
cat 1.6775 -0.4727 -0.9464 -0.2584
drank -1.6745 0.3439 0.3551 0.9755
milk 0.0804 0.4710 -1.6210 1.0696
5. feed_forward of that - each token on its own, no mixing
the 0.0887 -0.4677 0.1241 -0.6497
cat -1.2273 0.5467 0.6088 0.4716
drank -0.1238 -0.2006 -0.7904 -0.5756
milk -0.6148 0.0702 -0.2474 -0.1013
6. + the second residual: the block output
the 0.4402 -2.3647 1.3765 2.7099
cat -0.2339 0.6985 0.5752 0.7074
drank -2.1438 0.3589 -0.2166 0.7910
milk -0.3959 0.5558 -1.1902 0.7931
Only step 2 lets one token see another. Steps 4 to 6 treat every row separately,
which is why the feed-forward is described as position-wise and why two thirds of
a transformer's parameters do no mixing at all.
the same block written post-norm: ln(x + sublayer(x))
the 0.1304 -1.6706 0.7149 0.8252
cat 1.0187 -1.6334 0.5368 0.0779
drank -1.6229 0.3003 0.2135 1.1092
milk -0.0615 0.2700 -1.5009 1.2925
Stacking the same block over and over (weights reused, for a clean comparison):
depth pre-norm stream size post-norm stream size
0 1.1524 1.1524
1 1.2325 1.0000
2 1.4255 1.0000
4 2.4089 1.0000
8 4.9122 1.0000
12 7.8733 1.0000
pre-norm leaves the residual stream free to grow, so every large model ends with
one final layer_norm before the output head. post-norm normalises last, so the
stream is pinned - and the gradient has to pass through a normalisation on its
way back to every layer, which is why deep post-norm stacks need a warmup
schedule to train at all.
parameters in this block: 228
ok the block returns exactly the shape it was given
ok silence both sublayers and the pre-norm block hands back its input untouched
ok post-norm ends on a normalisation, so every output row has spread one
ok pre-norm does not, which is why a final norm sits before the output head
ok stacked twelve deep, the pre-norm stream has grown a long way
ok while the post-norm stream is pinned wherever the normalisation puts it
ok the block holds 228 parameters, 148 of them in the feed-forward
The two arrangements, and why you will see both:
- post-norm,
layer_norm(x + sublayer(x))— the original 2017 paper. Every output is normalised, so nothing can drift, but the gradient must pass through a normalisation on its way back through every layer, and the clean residual path is broken. Deep post-norm stacks are notoriously sensitive at the start of training, which is where learning-rate warmup came from. - pre-norm,
x + sublayer(layer_norm(x))— GPT-2 onwards, and essentially everything current. The residual path from input to output is now a pure sum with no normalisation on it, so the gradient route stays clean and very deep stacks train without heroics. The cost is the growth in the table above, which is why these models end with one finallayer_normbefore the output projection.
The order is not a stylistic preference and getting it wrong is a silent bug. Write
layer_norm(x) + sublayer(x) — normalise the residual branch instead of the
sublayer input — and the shapes are fine, the loss falls, and you have destroyed the identity
path the whole design rests on. Similarly, x + sublayer(x) with no norm anywhere
trains at small depth and diverges at large depth, for the reason the residual column in the
previous panel shows. When someone's hand-written transformer trains 10% worse than the
reference for no visible reason, this is the first thing to check.
Positional encoding, and the permutation that proves it is needed
Nothing in any panel so far knows what order the tokens arrived in. Attention scores every pair, and a set of pairs has no order. That is not a subtle point and it does not need to be taken on trust: permute the input and check what comes out.
# Positional encoding, and proof that without it the block cannot see word order.
import math
import random
TOKENS = ["the", "cat", "drank", "milk"]
D_MODEL = 8
def positional(position, d_model):
"""The 2017 sinusoidal encoding: alternating sine and cosine, at falling rates."""
out = []
for i in range(d_model // 2):
rate = 1.0 / (10000.0 ** (2.0 * i / d_model))
out.append(math.sin(position * rate))
out.append(math.cos(position * rate))
return out
PE = [positional(p, D_MODEL) for p in range(6)]
print("the encoding, one row per position, {} numbers each".format(D_MODEL))
print(" pos " + "".join("{:>9}".format("sin" if c % 2 == 0 else "cos")
for c in range(D_MODEL)))
for p, row in enumerate(PE):
print(" {:4} ".format(p) + "".join("{: 9.4f}".format(v) for v in row))
print()
print("Column 0 and 1 turn fastest, so they separate adjacent positions. The later")
print("columns turn slowly, so they carry roughly where in the sequence you are. No")
print("parameters at all: it is a fixed function of the position.")
def norm(v):
return math.sqrt(sum(x * x for x in v))
def dot(a, b):
return sum(x * y for x, y in zip(a, b))
print()
print(" every row has the same length:", ["{:.6f}".format(norm(r)) for r in PE])
print(" PE[0] . PE[k] as k grows :",
["{:.4f}".format(dot(PE[0], PE[k])) for k in range(6)])
print(" A position is most like itself and less like each of the next few. It then")
print(" turns back up, because sines repeat - the fastest column has a period of")
print(" about 6.3 positions. With d_model of 8 that wrap is visible at k = 4. In a")
print(" real model d_model is 768 and the slow columns keep distant positions apart.")
# ---- now the point. Attention on its own cannot tell one order from another ----
EMBEDDING = {"the": [0.2, -1.4, 0.9, 0.1, 0.5, -0.3, 0.0, 0.7],
"cat": [1.0, 0.1, -0.3, 0.4, -0.6, 0.2, 0.8, -0.1],
"drank": [-0.5, 0.7, 0.4, -1.2, 0.3, 0.9, -0.4, 0.2],
"milk": [0.3, 0.3, -1.1, 0.6, 0.1, -0.8, 0.5, 0.4]}
def matmul(rows, m):
return [[sum(r[k] * m[k][c] for k in range(len(m))) for c in range(len(m[0]))] for r in rows]
def softmax(row):
shift = max(row)
exps = [math.exp(v - shift) for v in row]
total = sum(exps)
return [e / total for e in exps]
def matrix(n, rng):
return [[rng.uniform(-0.5, 0.5) for _ in range(n)] for _ in range(n)]
rng = random.Random(42)
W_Q, W_K, W_V = matrix(D_MODEL, rng), matrix(D_MODEL, rng), matrix(D_MODEL, rng)
def block(rows):
"""Attention with no mask and no positions: exactly lesson 20's first panel."""
q, k, v = matmul(rows, W_Q), matmul(rows, W_K), matmul(rows, W_V)
scale = math.sqrt(D_MODEL)
out = []
for i in range(len(rows)):
w = softmax([dot(q[i], k[j]) / scale for j in range(len(rows))])
out.append([sum(w[j] * v[j][c] for j in range(len(rows))) for c in range(D_MODEL)])
return out
def summarise(rows):
"""A readout that does not care about order: average the positions."""
return [sum(r[c] for r in rows) / len(rows) for c in range(D_MODEL)]
ORDER = ["the", "cat", "drank", "milk"]
SHUFFLED = ["milk", "drank", "cat", "the"]
plain_a = block([EMBEDDING[t] for t in ORDER])
plain_b = block([EMBEDDING[t] for t in SHUFFLED])
print()
print("WITHOUT positions")
print(" sentence :", " ".join(ORDER))
for name, row in zip(ORDER, plain_a):
print(" {:8}".format(name) + "".join("{: 9.4f}".format(v) for v in row[:4]), " ...")
print(" reversed :", " ".join(SHUFFLED))
for name, row in zip(SHUFFLED, plain_b):
print(" {:8}".format(name) + "".join("{: 9.4f}".format(v) for v in row[:4]), " ...")
# The output rows came back in the new order, unchanged. Attention is permutation
# equivariant: permute the input and the output permutes with it, nothing else.
permuted_match = all(
close(a, b)
for t in ORDER
for a, b in zip(plain_a[ORDER.index(t)], plain_b[SHUFFLED.index(t)]))
print()
print(" every output row is bit for bit the same, just reordered:", permuted_match)
print(" so any order-blind readout gives the identical answer:")
print(" original :", ["{: .6f}".format(v) for v in summarise(plain_a)[:4]], "...")
print(" reversed :", ["{: .6f}".format(v) for v in summarise(plain_b)[:4]], "...")
# Add the position to the embedding. One addition, and the symmetry is gone.
placed_a = block([[e + p for e, p in zip(EMBEDDING[t], PE[i])] for i, t in enumerate(ORDER)])
placed_b = block([[e + p for e, p in zip(EMBEDDING[t], PE[i])] for i, t in enumerate(SHUFFLED)])
print()
print("WITH positions added to the embeddings")
print(" original :", ["{: .6f}".format(v) for v in summarise(placed_a)[:4]], "...")
print(" reversed :", ["{: .6f}".format(v) for v in summarise(placed_b)[:4]], "...")
gap = max(abs(a - b) for a, b in zip(summarise(placed_a), summarise(placed_b)))
print(" biggest difference between the two summaries: {:.6f}".format(gap))
print()
print("That is the whole of it. A transformer has no loops and no convolutions, so")
print("position is not implied by anything in the architecture. It has to be added to")
print("the input, and if you forget, the model trains happily on a bag of words.")
no_pos_gap = max(abs(a - b) for a, b in zip(summarise(plain_a), summarise(plain_b)))
test("every positional row has the same length, so no position is louder than another",
all(close(norm(r), norm(PE[0])) for r in PE), True)
test("the encoding is a function of position alone, with nothing learned",
positional(3, D_MODEL), PE[3])
test("a position is most like itself, and less like each of the next three",
[max(range(6), key=lambda k: dot(PE[0], PE[k])),
dot(PE[0], PE[1]) > dot(PE[0], PE[2]) > dot(PE[0], PE[3])], [0, True])
test("without positions, permuting the input just permutes the output", permuted_match, True)
test("so an order-blind summary of the sentence is exactly unchanged",
close(no_pos_gap, 0.0, 1e-12), True)
test("add the positions and the same permutation gives a different answer", gap > 0.05, True)What it printed, when the build ran it
the encoding, one row per position, 8 numbers each
pos sin cos sin cos sin cos sin cos
0 0.0000 1.0000 0.0000 1.0000 0.0000 1.0000 0.0000 1.0000
1 0.8415 0.5403 0.0998 0.9950 0.0100 1.0000 0.0010 1.0000
2 0.9093 -0.4161 0.1987 0.9801 0.0200 0.9998 0.0020 1.0000
3 0.1411 -0.9900 0.2955 0.9553 0.0300 0.9996 0.0030 1.0000
4 -0.7568 -0.6536 0.3894 0.9211 0.0400 0.9992 0.0040 1.0000
5 -0.9589 0.2837 0.4794 0.8776 0.0500 0.9988 0.0050 1.0000
Column 0 and 1 turn fastest, so they separate adjacent positions. The later
columns turn slowly, so they carry roughly where in the sequence you are. No
parameters at all: it is a fixed function of the position.
every row has the same length: ['2.000000', '2.000000', '2.000000', '2.000000', '2.000000', '2.000000']
PE[0] . PE[k] as k grows : ['4.0000', '3.5353', '2.5637', '1.9649', '2.2666', '3.1600']
A position is most like itself and less like each of the next few. It then
turns back up, because sines repeat - the fastest column has a period of
about 6.3 positions. With d_model of 8 that wrap is visible at k = 4. In a
real model d_model is 768 and the slow columns keep distant positions apart.
WITHOUT positions
sentence : the cat drank milk
the -0.1378 0.0571 0.1807 -0.2113 ...
cat -0.0420 0.0170 0.0549 -0.1542 ...
drank 0.0716 0.0664 -0.1648 -0.1836 ...
milk -0.0641 -0.1138 0.1436 0.0334 ...
reversed : milk drank cat the
milk -0.0641 -0.1138 0.1436 0.0334 ...
drank 0.0716 0.0664 -0.1648 -0.1836 ...
cat -0.0420 0.0170 0.0549 -0.1542 ...
the -0.1378 0.0571 0.1807 -0.2113 ...
every output row is bit for bit the same, just reordered: True
so any order-blind readout gives the identical answer:
original : ['-0.043068', ' 0.006697', ' 0.053570', '-0.128936'] ...
reversed : ['-0.043068', ' 0.006697', ' 0.053570', '-0.128936'] ...
WITH positions added to the embeddings
original : ['-0.281563', '-0.237581', ' 0.368353', '-1.054374'] ...
reversed : ['-0.347562', '-0.224593', ' 0.392750', '-1.016316'] ...
biggest difference between the two summaries: 0.091734
That is the whole of it. A transformer has no loops and no convolutions, so
position is not implied by anything in the architecture. It has to be added to
the input, and if you forget, the model trains happily on a bag of words.
ok every positional row has the same length, so no position is louder than another
ok the encoding is a function of position alone, with nothing learned
ok a position is most like itself, and less like each of the next three
ok without positions, permuting the input just permutes the output
ok so an order-blind summary of the sentence is exactly unchanged
ok add the positions and the same permutation gives a different answer
Three ways this is done, so job adverts stop being cryptic. Sinusoidal is
the panel above: fixed, free, no parameters, extends past the training length. Learned
absolute is an ordinary embedding table indexed by position, used by GPT-2 and BERT —
simple, and it hard-caps the context length at however many rows you trained.
Rotary (RoPE), which almost every current model uses, rotates
q and k by an angle proportional to position before the dot product,
so the score depends on the difference between two positions rather than their
absolute values. That relative property is what makes context-length extension tricks possible
at all.
The whole block, in a framework
Eight lines. Worth reading with the panels above in mind, because two of the arguments are choices this lesson spent panels on.
# The same block, as torch provides it.
# NOT RUN HERE - torch is not installed in this environment.
import torch
import torch.nn as nn
block = nn.TransformerEncoderLayer(
d_model=512,
nhead=8, # lesson 20's heads: 512 / 8 = 64 wide each
dim_feedforward=2048, # the 4 * d_model convention from the third panel
activation="gelu", # relu's smoother cousin, now the default in practice
norm_first=True, # TRUE is pre-norm. The default is False, i.e. post-norm.
batch_first=True,
)
# Written out, norm_first=True is exactly the third panel:
# h = x + attn(layer_norm_1(x))
# y = h + ff(layer_norm_2(h))
# and norm_first=False is:
# h = layer_norm_1(x + attn(x))
# y = layer_norm_2(h + ff(h))
model = nn.TransformerEncoder(block, num_layers=12)
final_norm = nn.LayerNorm(512) # the growth the third panel measured is why this exists
# Positions are NOT included. nn.TransformerEncoderLayer has no idea where a token is,
# so you add them yourself - and if you forget, nothing raises. This is the single most
# common bug in a hand-rolled transformer.
positions = nn.Embedding(max_len, 512) # learned absolute, GPT-2 style
x = token_embeddings + positions(torch.arange(seq_len))
# A causal decoder block wants the mask from lesson 20 as well:
mask = nn.Transformer.generate_square_subsequent_mask(seq_len)
out = model(x, mask=mask, is_causal=True)torch is not installed in this environment, so this panel was never executed and max_len, seq_len and token_embeddings are never defined in it. Both norm orders it describes were written out in plain Python above and their different behaviour under stacking was measured; the missing-positions warning is exactly what the fourth panel asserted.
Where this leaves you, and it is the end of the architecture. A transformer block is: normalise, attend, add; normalise, feed-forward, add. Attention is the only place tokens see each other — the feed-forward, which holds about two thirds of the parameters, works on each position alone. The residual is there so the gradient has an unattenuated path back to the input; the normalisation is there so the stream that residual builds does not run away. Position is not in the architecture and has to be added to the input. Stack that block twelve or ninety-six times, put an embedding table at the front and a projection to vocabulary logits at the back, feed those logits to lesson 6's softmax and lesson 19's sampler, and you have described every model this course is about.