Transformers

There is another problem with our model.

Currently it treats each position in the sequence separately. The same token gets different parameters depending on where it appears in the sequence.

That causes two problems:

• The model has to relearn each token at every position. If the model learned what United means when it appears at position 3, it still has to relearn it when United appears at a different position.
• The model has to relearn each relationship at every position pair. If the model learned that United followed by Kingdom at positions 3 and 4 is a country, it still has to relearn it for every other position pair they can appear in.

Attention

Both of these issues can be solved by the attention mechanism.

It introduces parameters that are shared across all tokens and lets every token look at every other token using the same parameters. So the token's meaning and its relationships are computed the same way wherever it appears.

Attention produces a contextual embedding for each token that carries information about all the other tokens in the sequence and its relationships to them.

Let's see how these contextual embeddings are calculated.

Key, Query and Value vectors

First we need to calculate three vectors for each token: Key (k), Query (q), and Value (v) by multiplying each token's embedding by shared parameters:

$$k_i = x_i W_K$$ $$q_i = x_i W_Q$$ $$v_i = x_i W_V$$

• $W_Q$, $W_K$, $W_V$ are the parameters that are shared across all tokens.
• The operation here is a matrix multiplication of vector $x$ by the $W_Q$, $W_K$, $W_V$ matrices

Attention score

Then for each token we look at every previous token in the sequence and decide how much attention to pay to each one. E.g. the relationship of ? to What's is probably more important than to others.

These values are called the attention score and they are calculated using this formula.

$$\text{score}_{ij} = \frac{q_i \cdot k_j}{\sqrt{d}}$$

• To get the score of $i$-th token attention to $j$-th token we multiply $i$-th token's Query by $j$-th token's Key.
• $i \geq j$
• $\sqrt{d}$ is needed to keep the score from being too large when $d$ is large.
• The multiplication is a dot product of two same-sized vectors returning a single number.

Blend the values

Now we can use all the scores for each token to get the contextual embeddings.

First, we turn the raw scores into attention weights by applying the softmax across each token's scores, so they become positive and add up to 1:

$$a_{ij} = \frac{e^{\text{score}_{ij}}}{\sum_{k=1}^{i} e^{\text{score}_{ik}}}$$

• $a_{ij}$ is the attention weight of the $i$-th token to the $j$-th token
• $i \geq j$

Then for each token we multiply every previous token's Value by its attention weight and sum them up.

$$\text{out}_i = \sum_{j=1}^{i} a_{ij} \cdot v_j$$

• The operation is a scalar multiplication of vector $v_j$ by the scalar $a_{ij}$

Here is the whole attention computation. I have decreased the context window to two tokens, to make it easier to visualise.

Note how the parameters are shared across all tokens and the relationships are learned per token pair.

Positional Encoding

We solved the problem with per-position parameters, but now we lost any information about the token's position. If we reorder the sequence, each token will get the exact same values of contextual embedding, but position still matters.

To fix this we will introduce another table (positional encoding) that carries values per position. The input of the attention for each token is a sum of the token's embedding and its positional encoding.

$$x_i = \text{embedding}_i​ + \text{position}_i​$$

Now the input carries both position and meaning at once. The positional encoding values are learned during training, the same way embeddings are learned.

Multi-head attention

There is also an option to run multiple attention heads ($H$) in parallel. Each head with its own $W_Q$, $W_K$, $W_V$ values.

At the end all heads' outputs are concatenated to a single contextual embedding per token.

This can help catch different kinds of relationships per attention head. But for our small LLM we will stick with $H=1$.

Transformers architecture

Once we have the contextual embeddings per token we can pass them to an MLP and get the next token predictions.

Here is how the complete transformer architecture for our LLM looks like with the context window increased to 32 tokens:

• The attention and MLP combine into a transformer block. If we want to scale it, we can stack multiple transformer blocks ($N$) on top of each other. But for our small model we will use only one block.
• Each token's MLP is independent as the contextual embeddings already have information about the sequence and can be computed in parallel.
• During training we split the corpus into context-window-sized blocks and pass each block to the transformer. Each token's predictions for that block will be calculated at once.

Here is how the transformer neural network trained on the same TinyStories corpus behaves for any input text:

We have built our small LLM. It might not behave exactly like popular frontier models do, but the main mechanism is the same.

If we want to improve its behaviour, the main thing needed is scaling: larger corpus, bigger embeddings, more parameters, etc. However, that scaling comes with a performance cost. Some frontier models use trillions of parameters, which makes them hard to run on one machine.

Even our small model has ~0.7 million parameters.

So we will leave the model as it is and implement it in code in the next section. Once we have the core architecture implementation, we can experiment more with the model and its hyperparameters.