Embeddings

The first thing we need to change is input representation. Using token IDs would not work for one reason.

It brings a wrong meaning to our model.

For example if we take our vocabulary the capital token has ID 1 and United has ID 7. When we pass these numbers as inputs, the model will do arithmetic with them. But the meaning of United is not bigger than capital.

These are just arbitrary numbers which have nothing to do with the actual meaning which makes it harder for the model to learn.

Embeddings table

The fix is instead of using random numbers we will use numbers that carry the meaning of the word. These numbers are called embeddings and they are vectors of size $d$.

The bigger the vector's size is, the more information each token can carry. For our vocabulary tokens let's start with embeddings of $d=2$.

If we plot each token's embedding on a 2D graph we can see how similar words are close to each other. The punctuation marks, places, helper words, and special tokens are grouped together as they have a similar meaning.

Now for any sequence we will pass each token's embedding instead of its ID. With four tokens in the context window the number of inputs increased to 8 ($d \cdot n$).

This is a much better representation but how exactly do we get the embedding values? These values are learned during the training, the same way the parameters are learned.

Training

The training process will not differ much from the spam filtering model besides a few things.

Embeddings update

First, we have to make sure we update the embeddings table after each run so the values become meaningful.

We do it the same way we update the parameters.

1. Calculate the derivatives with respect to each embedding value per token in the backward pass ($\frac{dL}{de}$)
2. Update the embedding values in the table using the derivatives ($e_{\text{new}} = e_{\text{old}} - \eta \cdot \frac{dL}{de}$)

Loss function

The loss function is the same cross-entropy we used for the spam filtering model.

As there is exactly one token that comes next for any training example, we calculate the loss for this token's output only (every other token's label in that training example is 0).

The formula always collapses to:

$$L = -\log_e(\hat{y}_{\text{correct}})$$

• $\hat{y}_{\text{correct}}$ is the probability our model gave to the token that actually came next.

Training set

Another difference is how we get the training set.

It is easier to do compared to the spam filtering model as we can generate the training set from one big text. We do not need humans to label each example.

Instead we take a big chunk of text (corpus) and break it into multiple input/output pairs.

For example, if our corpus is London is the capital of the United Kingdom., for four tokens in the context window we can generate this training set:

This type of training is called self-supervised learning.

Scaling the model

Now as we have improved our model with embeddings and know how to train it, let's scale to get better results.

First, let's increase our training set. We will use the public TinyStories corpus, which contains simple stories with about 3 million tokens.

If we search for the most common words that appear in that corpus, there are only 2000 tokens needed to cover 96% of the whole corpus. Let's increase our vocabulary to these 2000 tokens.

The neural network can be scaled as well, let's increase the context window to eight tokens, embedding size (d) to 128, and have 512 neurons in the hidden layer.

This is how the neural network looks like with these changes:

And here is how that model behaves after training: