Step 1 — A Bigram Model from Counts

Solution notebook. Run top to bottom; it is self-contained and takes under a minute. Read the prose — this notebook is written to be read, not just executed, and it is the starting point for Step 2.

It uses PyTorch throughout, the same library and the same idioms as the task page (torch.zeros, dim=/keepdim=, torch.multinomial, paired indexing), so everything here carries straight into Steps 2–8.

Checkpoint values quoted here were produced by this notebook, so your numbers should match to the digit (there is no randomness except in the sampling cell, which is seeded).

1.1 Load the data

On Colab, uncomment the download line. Locally, make sure input.txt is in the same folder as this notebook.

# !wget -q https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt

with open('input.txt') as f:
    text = f.read()

print(f'{len(text):,} characters')
print(text[:250])

1,115,394 characters — the complete works, as one string.

1.2 Tokenize

Our alphabet is whatever characters actually occur. stoi/itos are the two directions of a bijection $\{$characters$\} \leftrightarrow \{0,\dots,V-1\}$; encode/decode extend it to strings, and the assertion checks they are mutually inverse.

vocab = sorted(set(text))
V = len(vocab)
stoi = {ch: i for i, ch in enumerate(vocab)}
itos = {i: ch for i, ch in enumerate(vocab)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: ''.join(itos[i] for i in ids)

assert decode(encode('Fear no more')) == 'Fear no more'
print(f'V = {V}')
print(repr(''.join(vocab)))

V = 65. Token 0 is '\n' and token 1 is ' ' — the two most structurally important characters in the corpus sort to the front, which is a convenient accident. The vocabulary is an empirical fact about this text, not a design choice: all 26 letters appear in both cases, but the only digit is 3 (from a stray line number or two) and the punctuation is just ! $ & ' , - . : ; ? — no quotation marks, no parentheses. Editorial conventions of one 19th-century edition, fossilized into our model's alphabet.

1.3 Count

$N_{ab}$ = number of times character $b$ immediately follows $a$. The task page fills $N$ with a Python loop over zip(ids, ids[1:]), which walks the adjacent pairs one at a time — correct, and about ten seconds for a million pairs. Below is the vectorized version: flatten each pair $(a,b)$ to the single integer $aV+b$ and let torch.bincount tally all of them at once. (The loop is kept, commented out, so you can check the two agree.)

import torch
import matplotlib.pyplot as plt

ids = torch.tensor(encode(text))              # the whole text as a tensor of ids, shape (1115394,)
N = torch.bincount(ids[:-1] * V + ids[1:], minlength=V * V).view(V, V)

# The task page's loop gives exactly the same matrix, about 10 s slower:
# N_loop = torch.zeros((V, V), dtype=torch.int64)
# for a, b in zip(ids.tolist(), ids[1:].tolist()):
#     N_loop[a, b] += 1
# assert torch.equal(N, N_loop)

assert N.sum() == len(text) - 1
print(f'total pairs counted: {N.sum():,}')
print(f'cells that are exactly zero: {(N == 0).sum():,} of {V*V:,}')

2,822 of 4,225 cells are zero. Two thirds of conceivable character pairs never occur in all of Shakespeare. This is the concrete reason smoothing is not optional: an unsmoothed model assigns probability 0 to any of those pairs, and one occurrence in held-out text sends the log-likelihood to $-\infty$.

plt.figure(figsize=(9, 9))
plt.imshow(N.log1p(), cmap='Blues')
plt.xticks(range(V), vocab, fontsize=7)
plt.yticks(range(V), vocab, fontsize=7)
plt.xlabel('next character'); plt.ylabel('current character')
plt.title('log(1 + N)')
plt.colorbar(shrink=0.8);

We plot $\log(1+N)$ rather than $N$ because the raw counts span four orders of magnitude and a linear colour map would show only ' ' and 'e'. Three things to find in the picture, all verifiable from N:

• the q row is exactly rank one. q occurs 609 times and is followed by u all 609 times — probability 1.000, no exceptions in the complete works. English orthography as a single bright cell.
• the space column is bright but not universal: 46 of the 65 characters are ever followed by a space. Ask yourself which 19 are not, and you will have derived a chunk of English orthographic rules from a count matrix.
• the capital-letter block is brightest into itself — 51% of characters following a capital are themselves capitals, against only 33% lowercase. That is backwards from ordinary English prose, and it is the ALL-CAPS SPEAKER: convention of a printed play showing up as pure geometry. (A further 6.9% of post-capital characters are : itself.) The model has no concept of a play, a speaker, or a stage direction; it has a matrix in which those conventions are unmistakable.

q, u, sp = stoi['q'], stoi['u'], stoi[' ']
caps = torch.tensor([stoi[c] for c in 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'])
lower = torch.tensor([stoi[c] for c in 'abcdefghijklmnopqrstuvwxyz'])

print(f'q occurs {N[q].sum()} times, followed by u {N[q, u]} times')
print(f'{(N[:, sp] > 0).sum()} of {V} characters are ever followed by a space')
after_cap = N[caps].sum(dim=0)                       # what follows a capital, summed over capitals
print(f'after a capital: {after_cap[caps].sum() / after_cap.sum():.1%} capitals, '
      f'{after_cap[lower].sum() / after_cap.sum():.1%} lowercase, '
      f'{after_cap[stoi[":"]] / after_cap.sum():.1%} colon')

1.4 Normalize, with add-one smoothing

$P_{ab} = \dfrac{N_{ab}+1}{\sum_c (N_{ac}+1)}$, rows summing to 1. Laplace smoothing is the posterior mean under a $\mathrm{Dirichlet}(1,\dots,1)$ prior on each row — i.e. it is what you get by pretending you saw every pair once before looking at the data.

P = (N + 1).float()                 # counts are integers; probabilities need decimals
P = P / P.sum(dim=1, keepdim=True)  # divide each row by its own total

assert torch.allclose(P.sum(dim=1), torch.ones(V))
print('rows sum to 1 ✓')

The classic bug lives here. dim=1 sums along each row, and keepdim=True keeps the result shaped (V, 1) so it broadcasts down the rows. Drop keepdim and you get shape (V,), which PyTorch broadcasts along the last axis instead — silently normalizing columns. The assertion above is what catches it.

1.5 Sample

Run the Markov chain: start from the newline character and repeatedly draw the next character from the current row of $P$. torch.multinomial does the weighted draw; .item() turns the one-element tensor it returns back into a plain integer we can index with.

torch.manual_seed(0)
cur, out = stoi['\n'], []
for _ in range(500):
    cur = torch.multinomial(P[cur], num_samples=1).item()
    out.append(cur)
print(decode(out))

English-shaped nonsense: pronounceable syllables, plausible word lengths, capital letters after newlines, the occasional real word by chance. A single matrix of pair counts already captures that much. What it cannot capture is anything at a range beyond one character — which is the entire remaining agenda of the course.

1.6 Evaluate on the training text

Average log loss $\mathcal L = -\frac{1}{T-1}\sum_t \log P_{x_t x_{t+1}}$ in nats; divide by $\ln 2$ for bits per character; $e^{\mathcal L}$ is perplexity. The paired indexing P[ids[:-1], ids[1:]] pulls out all $T-1$ transition probabilities at once.

We wrap the computation in a function because Task 1.7 will run the very same pipeline on a different text.

def avg_log_loss(P, ids):
    p_next = P[ids[:-1], ids[1:]]     # entry t is P[ids[t], ids[t+1]]
    return -torch.log(p_next).mean()

def report(name, L):
    L = float(L)
    print(f'{name:<18} {L:.4f} nats  {L / torch.log(torch.tensor(2.)):.4f} bits/char  '
          f'perplexity {torch.exp(torch.tensor(L)):8.3f}')

P_uniform = torch.full((V, V), 1 / V)   # the baseline, run through the same code

report('bigram (add-one)', avg_log_loss(P, ids))
report('uniform', avg_log_loss(P_uniform, ids))

model	nats	bits/char	perplexity
uniform	4.1744	6.0224	65.000
bigram (add-one)	2.4549	3.5417	11.646

Perplexity is the effective branching factor: a uniform model over 65 characters is as hard to predict as a fair 65-sided die, and the bigram model reduces that to about 11.6. Shannon (1951) estimated English at roughly 1 bit/character, so at 3.54 we have captured perhaps a third of the available structure.

Start your table now — one row per step, and it becomes the story of the whole course:

step	model	bits/char
1	bigram counts	3.54

Three checks from the task page, each a one-liner:

# (a) the vectorized loss agrees with an explicit loop, on the first 10,000 pairs
loop = -sum(torch.log(P[a, b]) for a, b in zip(ids[:10000].tolist(), ids[1:10001].tolist())) / 10000
assert torch.isclose(loop, avg_log_loss(P, ids[:10001]))
print(f'loop vs vectorized on 10,000 pairs: {loop:.4f} vs {avg_log_loss(P, ids[:10001]):.4f}')

# (b) swapping the indices evaluates the model backwards — worse than knowing nothing
report('backwards', -torch.log(P[ids[1:], ids[:-1]]).mean())

# (c) the unsmoothed MLE is only 0.002 nats better on the training text
P_mle = N.float() / N.sum(dim=1, keepdim=True)
report('unsmoothed MLE', avg_log_loss(P_mle, ids))

1.7 Evaluate on held-out text

Everything above graded the model on the text its counts came from. Now split chronologically — first 90% training, last 10% validation — fit the counts on the training portion only, and score both. The vocabulary stays as built from the full text (conveniently, the last 10% introduces no new characters).

n_train = int(0.9 * len(ids))            # 1,003,854
ids_tr, ids_va = ids[:n_train], ids[n_train:]   # 111,540 validation characters
print(f'{len(ids_tr):,} train / {len(ids_va):,} validation')
print('validation opens:', repr(decode(ids_va[:45].tolist())))

N_tr = torch.bincount(ids_tr[:-1] * V + ids_tr[1:], minlength=V * V).view(V, V)
P_tr = (N_tr + 1).float()
P_tr = P_tr / P_tr.sum(dim=1, keepdim=True)

report('train', avg_log_loss(P_tr, ids_tr))
report('validation', avg_log_loss(P_tr, ids_va))
report('uniform (val)', avg_log_loss(P_uniform, ids_va))

2.4546 nats on train, 2.4819 on validation (3.5412 vs 3.5806 bits/char; perplexity 11.641 vs 11.964): a gap of 0.027 nats. The uniform model gives 4.1744 again, as it must — it doesn't depend on the data, so any change there is a bug.

The gap is real but tiny. The bigram matrix has only 4,225 cells sharing a million characters of evidence, so almost every cell is estimated from abundant data. Overfitting grows with the ratio of parameters to data, and here that ratio is small. Which transitions does the training portion never see?

train_count_of_val_pairs = N_tr[ids_va[:-1], ids_va[1:]]   # training count of each validation transition
unseen = train_count_of_val_pairs == 0
pairs = torch.stack([ids_va[:-1][unseen], ids_va[1:][unseen]], dim=1)
distinct, counts = torch.unique(pairs, dim=0, return_counts=True)

print(f'{unseen.sum()} validation transitions ({len(distinct)} distinct pairs) have zero training count')
for k in counts.argsort(descending=True)[:3]:
    a, b = distinct[k].tolist()
    print(f'  {itos[a]!r} -> {itos[b]!r}: {counts[k]} times')

P_tr_mle = N_tr.float() / N_tr.sum(dim=1, keepdim=True)
report('unsmoothed, train', avg_log_loss(P_tr_mle, ids_tr))
report('unsmoothed, val', avg_log_loss(P_tr_mle, ids_va))

187 validation transitions, 23 distinct pairs. The top three — S→P (63), E→B (42), N→Z (37) — are PROSPERO, SEBASTIAN, and GONZALO: The Tempest and its ALL-CAPS cast enter the corpus only in the final 10%. The unsmoothed MLE scores 2.4519 on train and infinite on validation; any one of the 187 suffices. That is the whole case for smoothing, in one number.

Going further — trigrams, and why tables must fail

Condition on two characters: $65^2 = 4{,}225$ contexts. We index the context pair $(a,b)$ as the single integer $aV + b$, and the count table becomes $4{,}225 \times 65$.

def trigram_P(ids):
    ctx = ids[:-2] * V + ids[1:-1]
    N3 = torch.bincount(ctx * V + ids[2:], minlength=V * V * V).view(V * V, V)
    P3 = (N3 + 1).float()
    return P3 / P3.sum(dim=1, keepdim=True), N3

def trigram_loss(P3, ids):
    ctx = ids[:-2] * V + ids[1:-1]
    return -torch.log(P3[ctx, ids[2:]]).mean()

P3, _ = trigram_P(ids)                      # fit and scored on the full text, like Task 1.6
report('trigram (full)', trigram_loss(P3, ids))

P3_tr, N3_tr = trigram_P(ids_tr)            # fit on the 90%, scored on both, like Task 1.7
report('trigram, train', trigram_loss(P3_tr, ids_tr))
report('trigram, val', trigram_loss(P3_tr, ids_va))
ctx_va = ids_va[:-2] * V + ids_va[1:-1]
print(f'{(N3_tr[ctx_va, ids_va[2:]] == 0).sum()} validation transitions unseen in training (bigram: 187)')

Trigram on the full text: 1.9532 nats = 2.818 bits/char, perplexity 7.05 — a large gain from one extra character of context. But that is a training-corpus number, and Lecture 1 §6.4 showed that $n$-gram training loss can only improve as context grows, all the way to memorization. The comparison that counts is the held-out one: refit on the 90%, the trigram scores 1.9526 train / 2.0684 validation (2.9841 bits/char, perplexity 7.91) against the bigram's 2.4819 — so it genuinely generalizes better here. The warning signs are growing, though: its train/validation gap is 0.116 nats to the bigram's 0.027, and 1,553 validation transitions are unseen in training, up from 187.

So why not keep going? A $k$-gram table has $65^{k+1}$ entries:

$k$	table entries	characters of data
1	$4.2\times10^{3}$	$1.1\times10^{6}$
2	$2.7\times10^{5}$	$1.1\times10^{6}$
3	$1.8\times10^{7}$	$1.1\times10^{6}$
5	$7.5\times10^{10}$	$1.1\times10^{6}$

By $k=3$ there are more table cells than characters in the corpus, so almost every cell is estimated from zero or one observation and the model is pure noise plus smoothing. The data does not grow; only the table does. This exponential blow-up against fixed data is the reason Lecture 2 replaces tables with parametrized functions, which can share statistical strength across contexts instead of treating each as an isolated counting problem.

→ Continue with Step 2, which rebuilds exactly this bigram model as $\operatorname{softmax}$ of a linear map and trains it by gradient descent.