Building Large Language Model from scratch

Building a basic implementation of a Large Language Model

· 8 min read

Background

You must have seen those movies where a hacker types bunch of words into terminal and a system is hacked. We all know hacking systems isn’t that easy. But to anyone outside of tech, it looks like magic - bunch of words go in, and magic happens. That is exactly how I felt when i first interacted with a LLM. It felt nothing like typical software. It was something entirely new. What is this LLM? How does it work? And what do people mean when they say LLM hallucinates?

The fastest way I know to get rid of those doubts is to build the thing myself. So that’s what this post is going to be about. I am going to build a language model from scratch in series of blog post. In the first blog, we will try to understand how an AI model behave and in subsequent ones, we will try to build a LLM and replicate the same behaviour.

First of all, what is LLM?

LLM (Large Language Model) is a neural networks that are designed to understand, generate and respond to human-like texts. In simpler terms, it is designed to predict next word based on the given texts. Before LLM, there were NLP (Natural Langauge Processing) which were useful in tasks of limited scope. It was used for tasks like email classification (classify a email as spam or not based on email text and sender) and sentiment analysis. LLM is just advanced and much powerful version of NLP.

Let’s look at the below figure. This is what a typical request and response look like in an LLM.

LLM takes some texts as input and returns a single word as response. The result i.e. “mat” is appended at the end of input text i.e. “The cat sat on the” and fed to LLM again in the next iteration. This goes on until LLM returns a special keyword <|endoftext|> when this loop comes out of the loop. In the below example, text on left side refers to input texts and text on the right side is result.

Iteration 1    →    "The cat sat on the"                       →    mat
Iteration 2    →    "The cat sat on the mat"                   →    and
Iteration 3    →    "The cat sat on the mat and"               →    fell
Iteration 4    →    "The cat sat on the mat and fell"          →    asleep
Iteration 5    →    "The cat sat on the mat and fell asleep"   →    <|endoftext|>

Building Blocks

Now we that know how a LLM generate texts, let’s go a bit deeper by splitting LLM into different blocks. In one iteration, the input goes through 4 stages and only two of them involves neural network.

Tokenizer

The transformer (the underlying architecture in an LLM) doesn’t really understand the English language as we do. To put it simply, a transformer is a mathematical model which understands only one datatype, i.e. numbers. So the tokenizer is the stage which converts our input text into a vector of numbers, which then gets passed to the transformer for further processing. The transformation happens in two stages.

The input text goes through two transformations before a model can make sense of it. In the first stage, the input text is converted into a vector of tokens. In the second stage, referring to the vocabulary, each token is assigned a unique ID. This results in a vector of token IDs, or numbers. Let’s play around with the visualization a bit and see how it works.

Input Text
Try:
Tokens
[
The cat sat on the mat and fell asleep
]
Token IDs
[
97690591013940229024503261815352547
]

We have seen how the text gets converted into token IDs. But wait, why do we even call them tokens ? Aren’t they just words ? Why not call them words and keep it simple? And what did I mean by vocabulary ? Let’s dig a bit deeper.

First, let’s focus on what a vocabulary is. It is nothing more than a lookup table — a hash map where each key is a piece of text and each value is the unique ID assigned to it: " cat" → 9059, " sat" → 10139. During the training phase, whenever the model encounters a new word, the model adds it to the map. As the model trains over enough data, the vocabulary gets built. In GPT 2 and GPT 3, there are 50,257 tokens and in gpt 4, there are 100,256 tokens in vocabulary.

Even with a vocabulary that large, a model will still meet text it has never encountered during training. Some models use a special token like <|unk|>. But every unknown word then collapses into same ID. GPT models instead use Byte Pair Encoding through which instead of using any special token, the unknown word is brokwn into subwords. Try Akwirw ier in the demo above and you get ['A', 'kw', 'ir', 'w', ' ier']. But as you can see, these are not valid words. So why are they in here? The vocabulary starts with the 256 possible byte values (“byte” in Byte Pair Encoding) and then it combines characters that frequently appear in english text. For example, the letter “e” appears after “d” in words like “define”, “made”, “hidden” and “delicate”.

Token Selection

Transformer includes tons of programming and maths. In layman’s terms, it looks at the input tokens and predicts next token. This stage is completely out of the user’s and even developer’s (the ones who work with agents with APIs) control. There is no way to tweak any parameter or do anything impacting model’s behaviour. It basically consumes input text and produces vector of next possible tokens with their probability percentages. This vector is then passed to this stage.

Token Selection stage is resposible for selecting one token out of many that are provided by the model. The way this stage chooses depends on several factors. First is the algorithm it uses. The most common way is greedy approach i.e. selecting token with the highest probability. It certainly works for deterministic problems i.e. maths or some programming problems where only a single answer is possible. In this case, greedy approach is our best bet.

But models are also used for creative work. They are tasked to write an essay or a story as well. If the model were to choose the most probable token, all essays or stories would sound the same. So in this particular scenario, instead of choosing the most probable token, we want the model to choose other tokens (with 2nd, 3rd or even 4th most probability) as well. Moreover, some parameters are exposed to tweak model behaviour:

  • Top-K: It is a number which tells the model to select next token among K tokens with highest probability.
  • Top-P: It is a percentage value which selects tokens (sorted by probability percentage in descending order) which have cumilative probability equal to or just over the Top-P value, and it selects a token from among them.
  • Temperature: This parameter takes a number between 0 and 2. It controls how random or predictable result will be.
    • Temperature being 0 lets model to choose the most probable token (the greedy approach we talked about)
    • Low temperature (0.1 - 0.3) makes the model’s result more deterministic. It is useful in case of code generation, math i.e. basically in contexts where only one answer is applicable.
    • High temperature (1.0 - 2.0) flattens the distribution making the model’s result more creative. It is useful in case of creative writing, brainstorming i.e. basically in contexts where we want different variations.

These three knobs are easier to feel than to read about, so here is the same distribution with all of them exposed. Pick a prompt, then drag the sliders and watch which tokens survive.

Prompt
Try:
Next Token keeping 49,840 of 50,257 tokens · top 30 shown cover 100.0%
#TokenProbabilityPCumulative
1 States
98.5%98.5%
2 Kingdom
1.3%99.8%
3 Arab
0.05%99.8%
4 Nations
0.04%99.9%
5 State
0.03%99.9%
6 S
0.02%99.9%
7.
<0.01%99.9%
8 East
<0.01%99.9%
9 Nation
<0.01%99.9%
10
<0.01%99.9%
11 U
<0.01%99.9%
12 United
<0.01%99.9%
13,
<0.01%99.9%
14 York
<0.01%100.0%
15 Republic
<0.01%100.0%
16 Empire
<0.01%100.0%
17 A
<0.01%100.0%
18 America
<0.01%100.0%
19 The
<0.01%100.0%
20 Auto
<0.01%100.0%
21 D
<0.01%100.0%
22-
<0.01%100.0%
23 St
<0.01%100.0%
24 Union
<0.01%100.0%
25 "
<0.01%100.0%
26 Way
<0.01%100.0%
27 City
<0.01%100.0%
28 We
<0.01%100.0%
29 Middle
<0.01%100.0%
30S
<0.01%100.0%