Crash course: from a line to a language model

Machine learning in one picture#

Classical programming means writing down the rules. Machine learning means showing examples and letting the computer find the rules. It always needs the same four things.

1. DataExamples of inputs and the right outputs. House sizes and prices, images and digits, text and the word that came next.
2. A modelA function with adjustable numbers in it, the parameters. Different parameters give different predictions.
3. A lossOne number that says how wrong the predictions are on the data. Smaller is better.
4. An optimiserA procedure that changes the parameters a little, again and again, so that the loss goes down.

The smallest example is a straight line. The model has two parameters, a slope $w$ and an intercept $b$:

The loss is the mean squared error, the average of the squared distances between the prediction and the truth:

And the optimiser is gradient descent. The gradient tells us, for each parameter, in which direction the loss goes up. We take a small step the other way. The size of the step is the learning rate $\eta$:

From a line to a neural network#

Real data is rarely a straight line. A neural network is a model that can bend. Its building block is the neuron: multiply the inputs by weights, add them up, and pass the sum through a simple non-linear function $\sigma$, for example one that sets negative values to zero.

Many neurons side by side make a layer. Layers stacked on top of each other make a network. Without the non-linear function, a stack of layers would still be one straight line. With it, a big enough network can approximate almost any function. That is why the same idea works for images, sound and text.

The gradient of the loss for every single weight is computed by the backward pass (backpropagation). You never do this by hand. PyTorch records the computation and does it for you. So the training loop of any model, from the line above to a GPT, is five lines:

for step in range(steps):
    x, y = get_batch()              # 1. a batch of examples
    logits, loss = model(x, y)      # 2. forward pass: predict and measure the loss
    optimizer.zero_grad()           # 3. forget the gradients of the last step
    loss.backward()                 # 4. backward pass: one gradient per parameter
    optimizer.step()                # 5. move every parameter a little

Predicting a class instead of a number#

A house price is a number. A handwritten digit is one of ten classes. For classes, the network outputs one raw score per class, called a logit. The softmax turns the scores into probabilities that are positive and add up to 1:

The loss for classification is the cross-entropy: the negative logarithm of the probability that the model gave to the right class $y$. If the model was sure and right, the loss is near 0. If it was sure and wrong, the loss is large.

This is the last piece of standard machine learning we need. Everything from here on is the same recipe, applied to text.

A language model is a next-token classifier#

The task every LLM solves is: given the text so far, guess what comes next.

Take the sentence that travels through this site: The Elbe flows through … You can finish it. So can a language model, and the way it does that is the whole story of this page.

The classes are all the pieces of text the model knows, its vocabulary. The model outputs one logit for each of them, the softmax makes a probability distribution, and the cross-entropy compares it with the piece that really came next. A sentence of ten tokens is ten classification problems.

The chain rule of probability says that this is enough to model whole texts. The probability of a text is the product of the probabilities of each token, given everything before it:

Nobody has to label any data for this. The text itself is the label: take any text, hide the next token, and ask the model. That is why language models can learn from enormous amounts of text.

Tokens: turning text into numbers#

A network only works with numbers, so the text has to be cut into pieces with integer IDs. There are three ways to cut.

Byte-pair encoding (BPE) starts from single bytes and keeps merging the most frequent pair of neighbours into a new token. Frequent words become one token, rare words are spelled out in pieces, and nothing is ever unknown.

A tokenizer is learned from text too, so it is good at the languages it saw most. That is a fairness question: you pay per token, in computing time, in money and in space in the context window. Notebook 6 measures this for six languages and three real tokenizers.

The simplest language model only counts#

Before any neural network, try the laziest idea: look only at the previous character and count what usually follows it. After q comes u. After t often comes h. This is the bigram model of notebook 4. Its text looks a bit like English letter by letter, but not word by word.

To compare language models we need a number. We use the cross-entropy on text the model has not seen, or its exponential, the perplexity:

You can read it as: the model is as unsure as if it had to pick uniformly from PPL options. A model that knows nothing about 65 characters has perplexity 65 and loss $\log 65\approx 4.17$. The bigram model reaches a loss of about 2.5. Its problem is the context of one character. It cannot know which word it is in.

Attention: letting tokens look back#

The fix is to let every position look at all the positions before it and collect what it needs. Not the positions after it, because those are what it has to predict. This rule is the causal mask.

Plain averaging treats every earlier token as equally important. Self-attention lets the weights depend on the content. Every token makes three vectors from its own embedding:

• a query: what am I looking for?
• a key: what do I have to offer?
• a value: what do I pass on if someone looks at me?

The score between two tokens is the dot product of one's query with the other's key. It is large when what one is looking for matches what the other has. A softmax over the scores gives the weights, and the output is the weighted average of the values. In one line, with the mask $M$:

In our sentence, the position after through has to find Elbe again, two tokens back. A river name there makes a city on that river likely. That is the kind of link attention can learn.

Nobody tells the model what to look for. The three small matrices that make queries, keys and values are parameters, and they are learned with the same loop as everything else.

The Transformer, and what GPT means#

One attention step is not enough, so the model does it several times, with a small ordinary network in between. One round of “talk to each other, then think alone” is a Transformer block. A GPT is a stack of them.

1. Token embeddingEvery token ID picks a row from a table of learned vectors. Which token is it?
2. Position embeddingA second learned vector is added for the position. Attention alone does not know the order of the tokens.
3. N Transformer blocksEach block is multi-head self-attention followed by a small MLP, with residual connections and layer normalisation so that a deep stack can still be trained.
4. LogitsA last linear layer gives one score for every token in the vocabulary, at every position.
5. Softmax and lossProbabilities for the next token, compared with the real next token by the cross-entropy.

Multi-head means several attention steps run side by side, each with its own queries, keys and values, so each can look for something different. GPT stands for generative pretrained Transformer: it generates text, it was trained on a lot of text beforehand, and it is built from these blocks.

In notebook 5 you write all of this in about 80 lines of PyTorch. The file train_tinygpt.py is the same model as a batch job.

Same recipe, more of everything#

The step from the model you train to the models in the news is not a new idea. It is more layers, wider layers, more text and more GPUs.

Generating text#

A trained model gives a probability for every possible next token. To write text we pick one, add it to the input, and ask again. How we pick is our decision, not the model's, and it changes the text a lot.

After The Elbe flows through, several continuations are right: Dresden, Hamburg, Magdeburg. Others sound just as natural and are wrong: Prague, Berlin. The model does not know the difference. It only has probabilities. In the demo below you decide how much risk it takes.

• Greedy: always take the most likely token. Always the same answer, and often repetitive.
• Temperature $T$: divide the logits by $T$ before the softmax. A small $T$ makes the model careful, a large one makes it adventurous and more often wrong.
• Top-k: keep only the $k$ most likely tokens.
• Top-p: keep the smallest set of tokens whose probabilities add up to $p$.

From a text predictor to a chatbot#

A model trained only to continue text is called a base model. Ask it a question and it may continue with another question, because that is a likely continuation. GPT-2 is a base model.

A chat model is a base model that was trained further on conversations. This is called instruction tuning. It expects its input in a fixed format, the chat template, where special tokens mark who is speaking. From the model's side a chat is still one long text, and answering is still predicting the next token, again and again, until it predicts the token that means “my turn is over”.

You see the raw template in section 9 of notebook 6, and generate.py in exercise 6 applies it to every prompt.

Where each idea lives in the material#