The reading assembled attention into full transformer blocks; here we watch the assembled machine learn to write. We train a tiny character-level transformer from scratch (two blocks, \(d = 64\), about 100K parameters) on roughly 2KB of text: the opening sentence of A Tale of Two Cities plus the Gettysburg Address. Training takes about half a minute on a plain CPU. Then we steer the model with prompts, and turn the temperature knob.
import timeimport torchimport torch.nn as nnimport torch.nn.functional as Fimport matplotlib.pyplot as pltimport scienceplotsplt.style.use(["science", "no-latex"])TEAL, CARDINAL, GRAY ="#009090", "#9c1b33", "#c9c9c9"torch.manual_seed(145)DICKENS = ("It was the best of times, it was the worst of times, it was the age of wisdom, ""it was the age of foolishness, it was the epoch of belief, it was the epoch of ""incredulity, it was the season of Light, it was the season of Darkness, it was ""the spring of hope, it was the winter of despair, we had everything before us, ""we had nothing before us, we were all going direct to Heaven, we were all going ""direct the other way - in short, the period was so far like the present period, ""that some of its noisiest authorities insisted on its being received, for good ""or for evil, in the superlative degree of comparison only. ")LINCOLN = ("Four score and seven years ago our fathers brought forth on this continent, a ""new nation, conceived in Liberty, and dedicated to the proposition that all men ""are created equal. Now we are engaged in a great civil war, testing whether that ""nation, or any nation so conceived and so dedicated, can long endure. We are met ""on a great battle-field of that war. We have come to dedicate a portion of that ""field, as a final resting place for those who here gave their lives that that ""nation might live. It is altogether fitting and proper that we should do this. ""But, in a larger sense, we can not dedicate - we can not consecrate - we can not ""hallow - this ground. The brave men, living and dead, who struggled here, have ""consecrated it, far above our poor power to add or detract. The world will ""little note, nor long remember what we say here, but it can never forget what ""they did here. It is for us the living, rather, to be dedicated here to the ""unfinished work which they who fought here have thus far so nobly advanced. It ""is rather for us to be here dedicated to the great task remaining before us - ""that from these honored dead we take increased devotion to that cause for which ""they gave the last full measure of devotion - that we here highly resolve that ""these dead shall not have died in vain - that this nation, under God, shall ""have a new birth of freedom - and that government of the people, by the people, ""for the people, shall not perish from the earth. ")corpus = DICKENS + LINCOLNvocab =sorted(set(corpus))m =len(vocab)stoi = {ch: i for i, ch inenumerate(vocab)}encode =lambda s: [stoi[c] for c in s]decode =lambda ids: "".join(vocab[i] for i in ids)print(f"corpus: {len(corpus)} characters, vocabulary: {m} distinct characters")print("vocabulary:", "".join(vocab))
This is exactly the reading’s block: \(h = 4\) heads with \(d_k = d/h = 16\), causal masking by setting future scores to \(-\infty\), and each sublayer wrapped in a residual connection and a layer normalization (callback: depth-enablers), with a \(d \to 4d \to d\) MLP. One piece here is not in the reading: we add a learned table of position vectors to the token embeddings. Without some position information, last lecture’s permutation equivariance means the model couldn’t tell “dog bites man” from “man bites dog”. Why that fix works, and a much better one, is next lecture’s entire subject.
The reading counted \(12d^2 + 4d = 49{,}408\) parameters per block; the model had better agree.
CONTEXT, D, HEADS, BLOCKS, D_FF =64, 64, 4, 2, 256class Block(nn.Module):def__init__(self):super().__init__()self.ln1, self.ln2 = nn.LayerNorm(D), nn.LayerNorm(D)self.Wq = nn.Linear(D, D, bias=False)self.Wk = nn.Linear(D, D, bias=False)self.Wv = nn.Linear(D, D, bias=False)self.Wo = nn.Linear(D, D, bias=False)self.mlp = nn.Sequential(nn.Linear(D, D_FF, bias=False), nn.ReLU(), nn.Linear(D_FF, D, bias=False))def forward(self, x): B, n, _ = x.shape dk = D // HEADS h =self.ln1(x) # pre-norm q =self.Wq(h).view(B, n, HEADS, dk).transpose(1, 2) # split into heads k =self.Wk(h).view(B, n, HEADS, dk).transpose(1, 2) v =self.Wv(h).view(B, n, HEADS, dk).transpose(1, 2) scores = q @ k.transpose(-2, -1) / dk**0.5# scaled dot products mask = torch.triu(torch.ones(n, n, dtype=torch.bool), diagonal=1) scores = scores.masked_fill(mask, float("-inf")) # causal: no peeking a = F.softmax(scores, dim=-1) out = (a @ v).transpose(1, 2).reshape(B, n, D) # concatenate heads x = x +self.Wo(out) # residual connection x = x +self.mlp(self.ln2(x)) # residual againreturn xclass TinyTransformer(nn.Module):def__init__(self):super().__init__()self.tok = nn.Embedding(m, D)self.pos = nn.Embedding(CONTEXT, D) # position info: next lecture's subjectself.blocks = nn.ModuleList([Block() for _ inrange(BLOCKS)])self.ln = nn.LayerNorm(D)self.head = nn.Linear(D, m, bias=False)def forward(self, idx): B, n = idx.shape x =self.tok(idx) +self.pos(torch.arange(n))for block inself.blocks: x = block(x)returnself.head(self.ln(x)) # logits, one row per positionmodel = TinyTransformer()print(f"parameters per block: {sum(p.numel() for p in model.blocks[0].parameters()):,}")print(f"total parameters: {sum(p.numel() for p in model.parameters()):,}")
parameters per block: 49,408
total parameters: 107,776
Training
Predicting the next character is classification over \(m = 37\) characters, so the loss is cross-entropy, exactly as in logistic regression. The causal mask makes all 64 positions of every sampled window valid, simultaneous training examples. A random guesser pays \(\ln 37 \approx 3.6\) nats, so the loss should start there and fall.
data = torch.tensor(encode(corpus))def get_batch(B=32): ix = torch.randint(len(data) - CONTEXT -1, (B,)) x = torch.stack([data[i:i + CONTEXT] for i in ix]) y = torch.stack([data[i +1:i + CONTEXT +1] for i in ix])return x, yopt = torch.optim.Adam(model.parameters(), lr=3e-3)start = time.time()for step inrange(1201): xb, yb = get_batch() loss = F.cross_entropy(model(xb).view(-1, m), yb.view(-1)) opt.zero_grad(); loss.backward(); opt.step()if step %200==0:print(f"step {step:5d}: loss {loss.item():.3f}")print(f"training time: {time.time() - start:.0f}s")
step 0: loss 3.789
step 200: loss 1.644
step 400: loss 0.286
step 600: loss 0.157
step 800: loss 0.120
step 1000: loss 0.118
step 1200: loss 0.100
training time: 27s
Steering with prompts
The weights are frozen from here on, so the only control left is the prompt. Three prompts, one model, temperature \(T = 0.8\): each prompt steers the model into a different groove of its training text.
def generate(prompt, n_chars=250, temperature=1.0): idx = torch.tensor([encode(prompt)])with torch.no_grad():for _ inrange(n_chars): logits = model(idx[:, -CONTEXT:])[0, -1] # logits for the next characterif temperature ==0: nxt = logits.argmax() # greedy decodingelse: p = F.softmax(logits / temperature, dim=-1) nxt = torch.multinomial(p, 1)[0] # sample idx = torch.cat([idx, nxt.view(1, 1)], dim=1)return decode(idx[0].tolist())for prompt in ["It was the ", "Four score", "The world"]:print(f'--- prompt: "{prompt}" ---')print(generate(prompt, 250, temperature=0.8), "\n")
--- prompt: "It was the " ---
It was the spring of hope, it was the winter of despair, we had everything before us, we had nothing before us, we were all going direct to Heaven, we were all going direct the other way - in short, the period was so herthing before us, work which the season, u
--- prompt: "Four score" ---
Four score and seven years ago our fathers brought forth on this continent, a new nation, conceived in Liberty, and dedicated to the proposition that all men are created equal. Now we are engaged in a great civil war, testing whether that nation, or any nation
--- prompt: "The world" ---
The world will little note, noredicred long end in a great cand long endure. We are met on a great battle-field of that war. We have come to dedicate a portion of that field, as a final resting place for those who here gave their lives that that case for rath
The temperature knob
We sample the next character from \(\operatorname{softmax}(\mathbf{z}/T)\). In class we proved the two limits: \(T \to 0\) recovers greedy argmax decoding, and \(T \to \infty\) the uniform distribution. Same prompt, four temperatures.
for T in [0.0, 0.5, 1.0, 2.0]:print(f"--- T = {T} ---")print(generate("It was the ", 200, temperature=T), "\n")
--- T = 0.0 ---
It was the spring of hope, it was the winter of despair, we had everything before us, we had nothing before us, we were all going direct to Heaven, we were all going direct the other way - in short, the period w
--- T = 0.5 ---
It was the age of wisdom, it was the age of foolishness, it was the epoch of incredulity, it was the season of Light, it was the season of Darkness, it was the spring of hope, it was the winter of despair, we ha
--- T = 1.0 ---
It was the epoch of incredulity, it was the season of Light, it was the season of Darkness, it was the spring of hope, it was the winter of despair, we holl enothis. Bed ded in vad styhomeaste a goitties the dea
--- T = 2.0 ---
It was the worst of times, it was the age of foolishness, it was car sonseceme a gil, birthorth liver, to ve -ut c. It is altogTo Heaven, we were all going direct tke the proposirtion that cont all goinent, yo g
Greedy repeats the training text verbatim (rerun the cell and \(T = 0\) gives the identical string), and \(T = 0.5\) stays close to it; \(T = 1\) starts inside the memorized text and then drifts out of it; \(T = 2\) dissolves into letter soup.
To see why, look at the distribution being sampled. The corpus contains “it was the” nine times, followed by four different characters: s three times (season, season, spring), w twice (worst, winter), a twice (age, age), and e twice (epoch, epoch). So the trained model’s next-character distribution is genuinely multimodal, and temperature reshapes it without reordering it.
with torch.no_grad(): logits = model(torch.tensor([encode("it was the ")]))[0, -1]chars = [c if c !=" "else"_"for c in vocab]top =int(logits.argmax())fig, axes = plt.subplots(3, 1, figsize=(7, 4.5), sharex=True, sharey=True)for ax, T inzip(axes, [0.5, 1.0, 2.0]): p = F.softmax(logits / T, dim=-1).numpy() ax.bar(range(m), p, color=[CARDINAL if i == top else TEAL for i inrange(m)], width=0.7) ax.text(0.99, 0.78, f"$T = {T}$", transform=ax.transAxes, ha="right") ax.set_ylabel("Probability")axes[0].set_ylim(0, 0.99)axes[-1].set_xticks(range(m))axes[-1].set_xticklabels(chars, fontsize=6)axes[-1].set_xlabel('Next character after "it was the " (the underscore is the space character)')plt.show()for T in [0.5, 1.0, 2.0]: p = F.softmax(logits / T, dim=-1)print(f"T = {T}: greedy pick '{vocab[top]}' holds probability {p[top]:.2f}")
T = 0.5: greedy pick 's' holds probability 0.79
T = 1.0: greedy pick 's' holds probability 0.51
T = 2.0: greedy pick 's' holds probability 0.28
The cardinal bar is the same character at every temperature, since dividing every logit by a common \(T\) never changes their order, so greedy’s choice is fixed. What changes is its share, which falls from \(0.79\) to \(0.28\); at \(T = 2\) probability leaks onto characters that never once followed “it was the” in the corpus, such as the b of “best”.
What if?
Change the prompt, the temperature, or the length. Good starting points: the prompt "we can not ", or temperature 5. (One caution: the tokenizer only knows the 37 characters printed at the top, so a prompt with any other character, like a question mark, will crash encode.)
prompt ="we can not "# change me!temperature =0.8# change me!n_chars =200# change me!print(generate(prompt, n_chars, temperature))
we can not consecrate - we can not hallow - this ground. The brave men, living and dead, who struggled here, have consecrated it, far above our poorthat we sh of herlishng pecrople, by the people, for the people
Punchline: a transformer is a next-token distribution plus a loop. The prompt picks where in its training distribution the model starts, the temperature picks how boldly it samples, and everything it “writes” comes, one character at a time, from a softmax we understand end to end.
