Next word prediction

When we prompt a request to an LLM it does not generate the whole output straight away. Instead it generates one unit of text (token) at a time.

So the LLM is called multiple times, extending the text after each call. This technique is called autoregressive generation.

We will build that model the same way we did with spam filtering: start with the simplest version first and gradually improve it.

Vocabulary

First, we need to create a vocabulary, the full list of tokens the model knows.

The vocabulary sets what tokens the model can recognise as input and what tokens it can produce as output.

Let's make a small vocabulary for our model. We will have each word that appeared in our example as a separate token plus three special tokens.

The last three entries are special tokens.

• <unk> is used for any other word that is not in the vocabulary.
• <eos> is used for the end of the sequence. When the sequence is complete, the LLM should return this token and we stop the generation loop.
• <pad> is used to fill empty slots in the input.

Now we can represent any input and output texts as numbers.

Input

For the input, we can assign and use a unique number for each token in the vocabulary. That way, What's the capital of the United Kingdom? sequence becomes 8, 6, 1, 5, 6, 7, 3, 10.

However that sequence can be very long depending on the input text. So we need to fix the number of tokens the model can receive in one call (this fixed window of tokens is called a context window).

• If the input text is longer than the context window, we drop the oldest tokens and keep only the most recent ones.
• If the input text is shorter than the context window, we pass the <pad> special token to fill the missing slots.

If we define the context window size as 4, we will always have four inputs (one for each token in the sequence in the same order). E.g. for the prompt What's the capital of the United Kingdom? we pass only the United Kingdom?

Output

The text generation is the same kind of classification problem as the spam filtering, but with more classes. Instead of choosing between spam or not spam, an LLM chooses between all the tokens in the vocabulary.

So our model will return a probability between 0 and 1 for each token in the vocabulary. That means we will have as many outputs as the vocabulary size.

For our vocabulary it is 13 outputs.

The output with the highest probability will be chosen as the next token.

Computation

For the computation we can reuse almost the same neural network we used for the spam filtering model, but with one change.

We can't use the sigmoid function ($y = \frac{1}{1 + e^{-z}}$) to get the final output as it will calculate the probability of each output in isolation and we can't compare them.

Instead we will use the softmax function.

• $n$ is the number of tokens in the vocabulary

This formula gives a valid probability distribution across all the tokens. Each result is a number between 0 and 1 and the sum of all results adds up to 1. If the probability of London is high, all the other probabilities must be lower.

Both sigmoid and softmax are called activation functions.

This is how our first LLM with four tokens in the context window, a 13-token vocabulary, and three neurons looks:

And here is how it behaves for any input text. Note that it is not trained at all for now, all the parameter values are random.