Teaching it to talk

Ask the Sprout from chapter 11 "Why is the sky blue?" and you won't get an answer. It may continue with another question, start a story, or write a dialogue between two strangers. It isn't being rude. Nobody ever told it that a piece of text can be a question addressed to it. It learned exactly one thing: how to continue documents. And to Sprout, a question is just the beginning of some document.

This chapter is about the smallest change that turns a text continuer into a conversation partner. We won't touch a single line of the model's code. Only the data changes: we show Sprout some twenty thousand documents of a special kind, conversations, and train it a little longer. This step is called supervised fine-tuning, or SFT, and it is how every chatbot you have ever used was made from its base model.

A storyteller doesn't answer

Below, the same question goes to two versions of Sprout: the base model from chapter 10 and the same model after fine-tuning on conversations. The base model can be asked in three different ways.

Try all three ways of asking the base model.

• Just the question. The base model treats it as the opening of a text and carries on: another line of dialogue, a narrator, sometimes a story. Nothing in its training said that a question at the start of a document must be followed by its answer.
• As a script. Here we wrap the question in a tiny scene: "Anna talks to her friend Sprout. Anna: … Sprout:". Sprout read plenty of such scripts in the SODA dialogues (chapter 9), and in a script the line after "Sprout:" is naturally a reply. Often this works surprisingly well. This trick, dressing the task up as a document whose natural continuation is the answer, was the main way to use language models before chat fine-tuning. GPT-3 (Brown et al., 2020) was used almost entirely this way.
• In the chat template. The question is wrapped in the same special marks the chat model gets. The base model has never seen them, so it is lost: it often starts in the middle of a word ("es can be bright…", "holders, I'm Sophia"), then drifts into a story or a dialogue between strangers that has little to do with the question.

The key observation is on the other side: the fine-tuned model answers. Yet it has exactly the same architecture, the same 17.31 million parameters, and was trained on only a small amount of extra text. The base model already knew grammar, how dialogues go, and the simple facts of a children's world. What it lacked was a format and a role: knowing that it is one side of a conversation and must answer the other side.

A script with special marks

The "script" trick hints at the solution: turn a conversation into a document of a fixed shape, and train the model on many such documents. We just need a shape that can't be confused with ordinary text. "Anna:" and "Sprout:" could appear in any story; we need marks that never occur in normal writing.

That is what special tokens are for. Back in chapter 6, when we built the tokenizer, we set aside four ids at the very end of the vocabulary: 8188 to 8191. Only one of them, <|endoftext|>, has been used so far, to separate documents in pre-training. Now the other three come into play:

• <|user|>: the person's turn starts here;
• <|assistant|>: Sprout's turn starts here;
• <|end|>: the current turn is over.

A whole conversation becomes one line of tokens:

<|endoftext|><|user|>hi!<|end|><|assistant|>Hello! How is your day going?<|end|><|user|>…

To chat, we build this line from the conversation so far, finish it with <|assistant|>, and let the model write until it produces <|end|>, the stop token promised in the previous chapter. Every chat box on this site does exactly that. Build a conversation of your own below and see how it turns into tokens:

What happens in a long conversation?

Every time you send a message, the whole conversation is rebuilt into one line and read again from the start: the model has no memory between turns other than this line. Sprout can see at most 512 tokens, and the reply needs room too, so our chat keeps only the most recent turns that fit into 440 tokens and quietly drops the older ones. That is why, in a long chat, Sprout forgets what you said at the beginning. Big chat models do exactly the same, only their window holds hundreds of thousands of tokens.

This is the function from sft.py that turns a conversation into training material. Besides the token ids it builds a second list of the same length, the mask: 1 where the model must learn to predict the token, 0 where it only reads.

def render(tok, messages): """Token ids and a 0/1 mask saying which positions the model should learn to predict.""" ids, mask = [tok['<|endoftext|>']], [0] for m in messages: role = tok['<|user|>'] if m['role'] == 'user' else tok['<|assistant|>'] body = tok.encode(m['content'].strip(), allow_special=False) + [tok['<|end|>']] ids += [role] + body learn = 1 if m['role'] == 'assistant' else 0 mask += [0] + [learn] * len(body) return ids, mask

Note allow_special=False. If someone types the characters "<|end|>" into the chat, the tokenizer encodes them as ordinary bytes, not as the special token. Otherwise a user could end Sprout's turn for it or pretend to be Sprout: a small-scale version of what security people call prompt injection.

Grading only the answers

In pre-training every token was a target: the loss was the average surprise over all positions. In fine-tuning we multiply each position's surprise by the mask and average only over the positions that count:

$$\mathcal{L} \;=\; \frac{\sum_t m_t \cdot \big(-\ln p_\theta(x_t \mid x_{\lt t})\big)}{\sum_t m_t}, \qquad m_t \in \{0, 1\}.$$

The model still reads every token: the user's words flow through attention and shape every prediction after them. The mask only decides where the error signal comes from. Why not grade the user's words too? Three reasons.

1. We want an assistant, not a mimic. Graded on the user's lines, Sprout would also learn to write like the users: "lol", typos, questions. Some of that would leak into its own turns.
2. The user's text is always given. When we chat, the person writes their own turns. Learning to predict them is effort spent on a skill that is never used.
3. User text is hard to predict. Nobody can guess what a person will type, especially at the start of a chat. Its surprise would dominate the average and drown out the signal we care about.

Two special tokens deserve a word. <|assistant|> is not graded: we insert it ourselves, so there is nothing to guess. The <|end|> that closes Sprout's turn is graded: that is how the model learns when to stop talking. Without it, Sprout would answer and then just keep going.

Here is the whole idea on a toy example. The probabilities are made up; watch how the user's word "sprout", which the model did not expect at all, drags the plain average up, and how the mask removes it.

import numpy as np # a toy conversation, one "token" per word; the mask marks what the model is graded on tokens = ["<|endoftext|>", "<|user|>", "hi", "sprout", "<|end|>", "<|assistant|>", "hello!", "how", "are", "you?", "<|end|>"] mask = np.array([0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1]) # the probability the model gave to each actual token (made up for the example) p = np.array([1.0, 0.9, 0.02, 0.001, 0.3, 0.95, 0.4, 0.6, 0.9, 0.8, 0.7]) surprise = -np.log(p) print("loss over every token: ", round(surprise[1:].mean(), 2)) print("loss over Sprout's turn:", round((surprise * mask).sum() / mask.sum(), 2)) print() # the first token is never predicted, so start from the second for tok, s, m in zip(tokens[1:], surprise[1:], mask[1:]): print(f"{tok:>14} {s:6.2f} {'graded' if m else '-'}")

In Sprout's model.py the same thing takes two lines at the end of forward: the cross-entropy is computed for every position with reduction='none', multiplied by the mask and divided by the number of ones.

Twenty thousand conversations

The template and the mask are the mechanism. What the model becomes depends on the conversations we feed it. Sprout's collection has 19,734 of them from four sources:

• 7,217 conversations written for Sprout's level across 18 topics: small talk, nature, science, daily life, stories, learning, friends, Sprout itself, the world, conversations that depend on earlier turns, technology, history, how-to, fun, health, the planet, simple explanations and a mixed bag. They follow one style guide: short sentences, words a ten-year-old knows, plain text without lists or emoji, and honesty about limits. Sprout says that it is small, that it only understands English, that it doesn't know today's date or the news, and that it can make mistakes. Each conversation appears twice in the mix, and the 400 about Sprout itself four times, so the model firmly learns who it is.
• 1,200 everyday conversations from the open everyday-conversations dataset: longer back-and-forth chats with an assistant.
• 2,500 story requests. The request names the hero of a TinyStories story ("Tell me a story about a little fox"), and the answer is that story, at most 170 words.
• 800 "continue this story" requests: the first sentence or two of a TinyStories story in the request, the rest in the answer.

Switch the strip to "by Sprout's tokens" and the picture changes. Written conversations are 77% of all conversations but only half of the graded tokens: they are short, 64 tokens on average, while a story request averages 178. Since the loss is an average over graded tokens, each source pulls on the model in proportion to its share of graded tokens, not of conversations.

That arithmetic caught us out. Our first attempt used 5,000 stories, up to 230 words long. The fine-tuned model was fluent and polite, but when someone wrote "I feel a bit sad today", it answered with a fairy tale. Of course: stories made up the biggest share of everything it had been graded on, so "answer with a story" had become its default. We halved the number of stories and kept only the short ones. Then came a second, subtler trap. The first story requests were built from SimpleStories' labels such as "bygone eras" or "an anti-hero", and the stories only loosely matched them. Asked for "a story about a brave little fox", the model cheerfully told one about a girl baking a cake: it had learned that a story need not fit the request. We rebuilt the requests from each story's own hero ("Once upon a time, there was a little fox…" becomes "Tell me a story about a little fox"), and the fox came back. (Ask "I feel a bit sad today" in the first widget to see how the final model takes it.) Data mixing is one of the least glamorous and most important decisions in building a language model.

Here is how the collection is assembled. conversation() checks that a chat starts with the user, ends with the assistant and alternates properly; story_requests() finds each story's hero in its first sentence and asks for exactly that story:

def conversation(messages, source): msgs = [{'role': m['role'], 'content': tidy(m['content'])} for m in messages if m['role'] in ('user', 'assistant')] ok = msgs and msgs[0]['role'] == 'user' and msgs[-1]['role'] == 'assistant' and all(m['content'] for m in msgs) ok = ok and all(a['role'] != b['role'] for a, b in zip(msgs, msgs[1:])) return {'messages': msgs, 'source': source} if ok else None STORY_ASKS = [ 'Tell me a story about {subject}.', 'Can you write a story about {subject}?', 'story about {subject} pls', 'I want a short story about {subject}.', 'Tell me a bedtime story about {subject}.', 'Write a story where the hero is {subject}.', 'Could you tell me a story about {subject}, please?', 'Make up a story about {subject}!', ] # "Once upon a time, there was a little fox named ..." -> "a little fox" SUBJECT = re.compile(r"\bthere (?:was|lived) (a|an) ([a-z]+(?: [a-z]+){0,3}?)(?= named| who| that| in | with| called|,|\.)") def story_requests(raw, n, rng, per_noun=120): """Ask for a story about exactly the hero the story is about, so request and answer agree. (A first version asked for SimpleStories' abstract topic labels like "bygone eras"; the stories only loosely matched them, and the model learned that a story need not fit the request.) """ with open(f'{raw}/TinyStoriesV2-GPT4-train.txt', encoding='utf-8') as f: text = f.read(150_000_000) per = {} out = [] stories = text.split('<|endoftext|>')[1:-1] rng.shuffle(stories) for s in stories: s = tidy(s) m = SUBJECT.search(s[:220]) if not m or len(s.split()) > 170 or not 1 <= len(m.group(2).split()) <= 3: continue noun = m.group(2).split()[-1] if per.get(noun, 0) >= per_noun: continue per[noun] = per.get(noun, 0) + 1 ask = rng.choice(STORY_ASKS).format(subject=f'{m.group(1)} {m.group(2)}') out.append(conversation([{'role': 'user', 'content': ask}, {'role': 'assistant', 'content': s}], 'story')) if len(out) >= n: break return out

Altogether, the conversations come to 1.8 million tokens, and 75% of them belong to Sprout's turns. Compare that with the 330 million tokens of pre-training: the fine-tuning data is about half a percent of what the model has already read.

Fine-tuning itself

sft.py is a close cousin of train.py from chapter 10. It starts from the base model's weights instead of random ones; otherwise it is the same model, the same optimisers (Muon for the matrices, AdamW for the rest) and the same forward pass, now with the mask. What differs is the scale and the care:

• The learning rate is five times lower than in pre-training (these are the defaults in sft.py): 0.004 instead of 0.02 for the matrices, 0.0008 instead of 0.004 for the embeddings. The model already knows a lot, and big steps would trample it.
• By default the data is shown three times (three epochs) in batches of 32 rows, with 20 warm-up steps and a linear fade of the learning rate over the second half.
• We hold out 3% of the conversations. They are never trained on, only used to measure the loss, so we can see whether the model is learning to talk or just memorising.

Conversations have different lengths, and the model trains on rows of fixed length. pack() glues conversations one after another into rows of exactly 513 tokens: 512 inputs plus one extra, because the targets are the same row shifted by one position. Whatever doesn't fit starts a new row, and the leftover space is filled with padding, which the mask ignores.

def pack(examples, context): """Glue conversations into rows of exactly context+1 tokens (the rest is padding).""" rows, cur_ids, cur_mask = [], [], [] for ids, mask in examples: if len(ids) > context + 1: continue if len(cur_ids) + len(ids) > context + 1: rows.append((cur_ids, cur_mask)) cur_ids, cur_mask = [], [] cur_ids += ids cur_mask += mask if cur_ids: rows.append((cur_ids, cur_mask)) x = torch.zeros(len(rows), context + 1, dtype=torch.long) m = torch.zeros(len(rows), context + 1, dtype=torch.long) for i, (ids, mask) in enumerate(rows): x[i, :len(ids)] = torch.tensor(ids) m[i, :len(mask)] = torch.tensor(mask) return x, m

Here is the same packing done in your browser on the sample conversations from the explorer above:

Sprout's conversations are short: the median is 66 tokens and the longest is 359, so none has to be thrown away. The training part packs into about 3,900 rows that are nearly 90% full; with the default settings that makes about 370 steps. Next to the 10,070 steps of pre-training, that is a light touch.

How many epochs? Too few and the new format doesn't stick; too many and the model starts memorising particular conversations instead of learning how to converse. The held-out conversations act as referee. While their loss goes down, the model is generalising. If it turns up while the training loss keeps falling, the model is memorising. Here is Sprout's real fine-tuning run:

What changed inside

The loss curve is an average. To see the change up close, let's give both models the same conversation, one they have never seen, and measure how surprised each is by every token of Sprout's reply. This is exactly the quantity fine-tuning pushes down.

Two things stand out. The ordinary words of the reply become less surprising, especially the first ones: the base model doesn't expect an answer to begin right after <|assistant|>, because it has never seen that token. And the final <|end|> changes most of all. For the base model it is a token that never once appeared in its training data, so its probability is vanishingly small. After fine-tuning, Sprout expects to close its turn at the right moment: in the first two conversations it gives <|end|> more than 99%. The frog is the exception worth a look. After a story of just two sentences the chat model still expects more (the stories it was trained on went on for a paragraph or two), so <|end|> gets only a few chances in a hundred thousand. Even that is a thousand times more than the base model gives it.

Why the base model thinks <|end|> is nearly impossible

Recall from chapter 3 how softmax and cross-entropy push on the logits: the gradient of the loss with respect to the logit $z_j$ is $p_j - y_j$, where $y_j$ is 1 for the correct token and 0 for all others. For a token that never once appears in the training data, $y_j$ is always 0, so its gradient $p_j$ is always positive, and every step of gradient descent lowers its logit a little. Over ten thousand steps of pre-training, <|user|>, <|assistant|> and <|end|> were pushed down at every position of every batch and never pulled up once. Since Sprout's input and output embeddings are shared (chapter 8), the same rows also describe these tokens when they appear in the input, which is why the base model gets so lost when it reads them. Fine-tuning is the first time anything pulls them up.

The weights themselves moved only a little: a few hundred steps at a low learning rate. But the behaviour changed completely. That is the pattern of fine-tuning in general: a small nudge in weight space, a big change in what the model does.

Why is the <|end|> that closes Sprout's turn graded, while <|assistant|> is not?

<|assistant|> is too rare to learn We insert <|assistant|> ourselves, but only the model can decide when its answer is finished Special tokens are never graded; <|end|> is an exception for technical reasons

When we chat, our code appends <|assistant|> to hand the turn to Sprout, so predicting it would be pointless. The end of the answer is Sprout's own decision, and if it never learned to make it, it would never stop.

What fine-tuning can't do

It's tempting to think fine-tuning teaches the model what it says. Mostly it doesn't. It teaches the form: to answer rather than continue, to be brief, to be kind, to stop. The content of the answers comes almost entirely from pre-training. The LIMA paper (Zhou et al., 2023) made the point sharply: a large base model fine-tuned on just 1,000 carefully chosen conversations was surprisingly often judged as good as, or better than, far more heavily tuned models. The authors called it the superficial alignment hypothesis: nearly all knowledge is learned in pre-training, and fine-tuning mainly selects the style in which to use it.

The flip side is that fine-tuning can't add knowledge the base model doesn't have. Ask Sprout "What is 17 times 23?" in the first widget. The chat version answers at once, with no trace of doubt and in the shape of an answer, "this is that", because answering is what it was trained to do. But the numbers inside are a jumble: in a dozen tries we got "7 times 3 is 45", "19 = 56" and "980 is 7 times 36", and never 391. Nothing in children's stories taught it multiplication; even for "2 plus 2" its 4 usually comes wrapped in nonsense like "30 is 4". This is the heart of what people call hallucination: the model has learned the form of a helpful answer perfectly and fills it with whatever its knowledge suggests. Our data teaches Sprout to be honest about the limits we could predict: no date, no news, English only, shaky at anything beyond simple maths. But no data can teach it to notice every gap in its own knowledge. Gudibande et al. (2023) saw the same thing at a larger scale: small models fine-tuned on a big model's answers copied its style convincingly, and much less of its accuracy.

Sprout right now

This is the Sprout we have been growing all along: the same 17.31 million parameters, now fine-tuned to hold a conversation. Say hello, ask it about animals or the weather, ask for a story, tell it how your day went. It is small, so it sometimes gets facts wrong or loses the thread in long conversations, and it only understands English. In the next chapter you will learn to fine-tune it yourself, giving it a style of its own without retraining all 17 million numbers.