How LLMs Turn Text into Embeddings

This is part of my learning to build LLMs from scratch. I am following Sebastian Raschka’s Build a Large Language Model (From Scratch) as I work through each stage.

The workflow for training LLMs is something like this:

text -> tokenised text -> token IDs -> shifted X/Y sequence -> token embeddings -> add positional embeddings -> transformer model -> vocabulary logits -> compare with Y -> loss -> backpropagation

GPT-style LLMs are trained to predict the next token given the preceding tokens. For large language models to work with text, they need to be converted to numeric form. So something like “my best food is rice and beans” needs to somehow map to [1820, 1266, ...]

A naive way to implement this mapping is to split the sentence above into words, normalise each word and then map each word to a predefined number. This works, but there are two important nuances:

1. Building a map for every single word would be huge. Merriam-Webster’s dictionary estimates that there are over 470,000 entries.
2. Handling unknown or misspelt words would be tricky

This is where Byte pair encoding (BPE) comes in. Engineers can access this tokeniser through the tiktoken library. It works by breaking words into subwords, so unhappy can become something like ["un", "happy"]. The GPT-2 encoding available through tiktoken has a vocabulary of 50,257 tokens. These tokens can represent complete words, parts of words, punctuation or byte sequences.

The algorithm ensures that arbitrary text can be tokenised (broken into smaller words) using a subset of this vocabulary. For instance, I invented a word called “ghanfhete” which returns four tokens: [6064, 69, 258, 660] which maps to ['ghan', 'f', 'he', 'te']. This solves the first problem of trying to capture all possible word variations and representing words even when the complete word was never present in the tokeniser’s training data.

BPE handles both tokenising the text and mapping to token IDs.

From our graph above, the next stage is creating embeddings for each token. A good question to ask is if tiktoken already returns IDs which are numeric, why do we need embeddings? The answer is because the IDs don’t encode any meaning at all. Their only purpose is to map a token to a number. An LLM needs to “learn” what each number means in the context of other numbers. During training, the token embeddings develop useful base representations. The Transformer then produces context-dependent representations from them. To start with, we assign each token to an embedding with random initial values. The embedding matrix has shape [vocabulary_size, embedding_dimension]. If we use the GPT-2 encoding with an embedding dimension of 256, our embedding matrix has shape torch.Size([50257, 256]).

Remember those token IDs from tiktoken for each token? Yep, we can then use them to access the embedding for that token. For example, since "ghan" has token ID 6064, we retrieve embedding_matrix[6064]. This is row index 6064, or the 6,065th row when counting normally.

These embeddings are updated during training so that they become useful starting representations for their tokens. The Transformer later combines them with the surrounding tokens to produce context-dependent representations. More on that when I understand it myself, lol.

Finally, we have embeddings for each token; what about positional embeddings? why do we need them? The way GPT-style models work is that each token attends to itself and the tokens preceding it. Attention performs mathematical operations that combine information from these tokens. However, the token embedding itself does not record where the token appears in the sequence. To give the model this information, we create a separate embedding for each possible position: position 0 gets one vector, position 1 gets another, and so on. Because the positional and token embeddings have the same length, we can add them together:

token embedding + positional embedding = input embedding.

That’s the gist at a high level. Next week, we’ll explore coding attention mechanisms.