Creating an LLM (little language model) from scratch — Part I
Along Came the Water
ACT I
SCENE I
A lonely tavern beside a swollen river. Rain beats against the windows. Ophelia sits beside the fire. Enter Macbeth, carrying a broken lantern.
OPHELIA
The night grows cold; why dost thou tremble so?MACBETH
The cold lies not within the air, but here,
Where guilty thoughts do knock against my heart.OPHELIA
What thought could make so brave a soldier pale?OPHELIA
Blush, thou gentlewoman, and loved.
Thou shalt have as good as the devil wait on thee.MACBETH
I have heard thy name of Macbeth: I dare not;
Two doors are like fresh and blood, and thou leaves me:
As sailors do see their most rule moan. Come, my Lord Master Richard,
I will teach you the gown with a moon-born man,
And carry me this foot of all vile honour in my soul
And so willing relieved thy demands.MORTIMER
I had rather be my son to that billiard,
Himself would have wings for a king.OPHELIA
How fares you?MACBETH
Who bears his throat? He looks too deep.
His master’s tongue shake his throat up, and all miserable
Strike him forth his death. But he’s mad at me! I have seen him.MACBETH
I do weary, and think you still slay
For he is now in shape that he does so;
’Tis strange:
And when you shall be won by him; what he say to love,
Let me be so hungry: and what is he to?MALCOLM
Send you there, Macbeth!OPHELIA
That he is my father’s son.MACBETH
What art thou, Master Richard says.MACBETH
Why, that I am glad on’t. When the old man’s a bragged maid woos are
out. — Go away before him! Take heed, thou fool,
About thy soul!MACBETH
I think thou wilt speak nothing. I am thy son;
Thou canst not speak of us but tomorrow morning.
This day shall befall thee home; my wife is
These days, we are fascinated (scarred, annoyed, <insert state of mind>) by the capacities of LLMs (Large Language Models). With each new model, each new generation meets a new frontier of what can be achieved. Generative machinery is no longer a mirage; it is real and at the tip of our fingers (and voice).
However, it is fun to go back to basics and understand what this is all about. What is the sorcery behind the generative capabilities of Codex or Claude? That’s why, as a pet project, I decided to build a little language model from scratch with a Shakespearean vibe. And because everyone and their mother is using Python, I decided to use Swing ;)
What is an LLM after all?
By now, given the gazillion articles online dedicated to LLMs, the idea of what an LLM is should be generalised, but, for completion, I’ll state it: an LLM is a very elaborate probabilistic machine.
Given a sentence, e.g. “Pizza should not have”, an LLM will compute what the most likely word after “have”. As an example, ChatGPT replied:
“Pizza should not have… an empty spot where the cheese was supposed to be.”
while Claude replied:
“Pizza should not have… too many toppings competing for the same square inch.”
(because ChatGPT and Claude are tuned to follow instructions, I had to prompt them: “Complete the sentence: ‘Pizza should not have’”)
Now, the LLM will not generate the whole sentence in a single sweep. Given “Pizza should not have”, an LLM will generate
“Pizza should not have… “
“Pizza should not hot have… an”
“Pizza should not have… an empty”
“Pizza should not have… an empty spot”
and so forth, until an end-of-file is emitted or the context is exhausted.
As you can see, the LLM predicts one word at a time. (Now you know why you have to wait so long for complex responses) and each word it emits is the one with the highest probability.
How the heck does it know which is the most likely word that comes next?
That is a very fine question, and the raison d’etre of training. But first, please pay attention to the following.
Get Bruno Oliveira’s stories in your inbox
Join Medium for free to get updates from this writer.
Artificial Intelligence is not something new. It has been around for ages, and many methods have been developed to create “thinking machines”, especially in NLP (Natural Language Processing). However, no matter the algorithms and techniques (neural networks, deep learning), the results were less than satisfactory. (Many of you still remember the first versions of Google Translate).
The main reason for the failure was that there was no concept of relationship, that is, how words interconnect with each other, and how concepts are related to each other.
A simple example:
“Hoje é um bom dia.” is Portuguese for “Today is a good day.”. We can effectively translate word by word, and both sentences will be correct.
However, “Hoje é um bom dia?” is now a question, and so it translates to “Is today a good day?” and we can no longer translate the sentence word by word.
Of course, this is a simple example, and it would be easy enough to break the original sentence, find the question mark, interpret it as a question, and apply the English rules for questions. But the beauty of an LLM is that it does not need to know the rules because they will emerge during training.
How? Because of the simplicity of something called attention.
Just like the invention of Calculus opened the door to a whole new branch of mathematics (and loads of headaches for students), the attention mechanism (developed in 2014 by Bahdanau et al. and later used by Vaswani et al. to create the transformer architecture, in the most famous paper in the history of AI) came to shake and revolutionise the AI world.
In the words of Morpheus
We marvelled at our own magnificence as we gave birth to AI.
So how does this attention thing work?
To try to understand attention and its importance, let’s look at an example. Imagine that you’re given a topographic map of Great Britain and tasked with driving from London to Edinburgh.
The only notion you have is that one road connects to others at intersections, and so your options are to start on a road and randomly choose another every time you reach a crossroad. Eventually, you will end up somewhere, say in Dover.
Dover is not Edinburgh, so you did not reach your final destination. What do you do? You do what neural networks do in training: you backtrack. You register the roads that took you to Dover, and then you move some roads back (I’m assuming that all roads are two-way), and you pick at some intersection a road not travelled, which will take you to another destination.
As you can now figure it out, this method of backtracking will hardly take you to Edinburgh, and the reason why is that you only have local knowledge. By only knowing the road connections at each intersection, you cannot effectively decide which road will take you to your destination. You can rely on the information you already gathered (all other attempts), but apart from that, you are using mostly random choices.
Now, imagine that you have a little more information. You know that Edinburgh is to the northeast of London, and that, for instance, Edinburgh is closer to Manchester than it is to Birmingham, and that Birmingham is south of Manchester. The next time you pick a road, you can use that information to weigh your decisions. You will be more likely to pick a road that heads northeast, rather than the one that heads west.
In the same vein, if you reach Birmingham, you know you are still far and that you should not go south, because Manchester is closer to Edinburgh and you are south of it.
Again, once you reach a destination and need to backtrack, you will weigh your knowledge to choose another road. Now it will not be entirely random, because you have more than just local information; you have spatial relationship information. And this last part is what attention offers.
From roads to words to tokens
The attention mechanism is powerful because it provides relations. Depending on the number of heads (more on this soon), the model will learn that a question mark is a question, that the correct order of adjectives is opinion, size, age, shape, colour, origin, material, and purpose and that a cat is closer to a tiger than a spider. Structure and related topics all emerge from attention.
Because LLMs are language models, we need language, which is to say words. However, annoyingly, computers would rather work with numbers than with orders. And this is why, in LLMs, we use tokens instead of words.
A token is a number that represents part of a word (sometimes a whole word, like the or and), and they are built based on a corpus (a corpus is any amount of text we provide an LLM in training).
To break words into tokens, we can start with the most basic representation, ASCII, where each letter and symbol is associated with a number between 0 and 254. As an example, the sentence Hello, World! would have been tokenised as [72, 101, 108, 108, 111, 44, 32, 87, 111, 114, 108, 100, 33]. However, this representation would be poor for our intent, because we would need to start by inferring connections between letters, and only then connections between words.
A better solution is to perform a first-word-to-ASCII conversion, then search for repeated sequence clusters. Let’s see an example. Take the sentence
playing players played playfully on the playground while replaying the gameplay.
The first step will be to take each character and create its ASCII representation, which would be something like
112, 108, 97, 121, 105, 110, 103, 32, 112, 108, 97, 121, 101, 114, 115, 32, 112, 108, 97, 121, 101, 100, 32, 112, 108, 97, 121, 102, 117, 108, 108, 121, 32, 111, 110, 32, 116, 104, 101, 32, 112, 108, 97, 121, 103, 114, 111, 117, 110, 100, 32, 119, 104, 105, 108, 101, 32, 114, 101, 112, 108, 97, 121, 105, 110, 103, 32, 116, 104, 101, 32, 103, 97, 109, 101, 112, 108, 97, 121, 46
With this list, we then start searching for candidate pairs, that is, consecutive numbers that appear many times.
A good candidate pair is 112, 108, corresponding to “pl”, with seven occurrences. Having a good candidate, we assign a new number corresponding to that pair (usually 255, the next number after the complete ASCII table).
255, 97, 121, 105, 110, 103, 32, 255, 97, 121, 101, 114, 115, 32, 255, 97, 121, 101, 100, 32, 255, 97, 121, 102, 117, 108, 108, 121, 32, 111, 110, 32, 116, 104, 101, 32, 255, 97, 121, 103, 114, 111, 117, 110, 100, 32, 119, 104, 105, 108, 101, 32, 114, 101, 255, 97, 121, 105, 110, 103, 32, 116, 104, 101, 32, 103, 97, 109, 101, 255, 97, 121, 46
And we continue this process, creating new numbers for known pairs. (Mind that now 255, 97 is a pair, corresponding to “pl” and “a”).
This process may find tokens such as play, ing, ed, er, ers, and it will stop when no more candidate pairs exist, or we have exhausted the number of allowed tokens.
The algorithm I just described is called Byte-pair encoding (BPE) and, in some form or another, is used in LLMs to transform words into text and can be quickly implemented in a naive way.
idToToken = (0..<256).map { [UInt8($0)] }
tokenToId = Dictionary(uniqueKeysWithValues: idToToken.enumerated().map { ($0.element, $0.offset) })
let ordinaryMaximum = maximumVocabularySize - definitions.count
var tokens = byteTokens(from: text)
while idToToken.count < ordinaryMaximum {
updateLine("Token count: \(idToToken.count)/\(ordinaryMaximum)")
guard let winner = mostFrequentPair(in: tokens) else { break }
merges.append(winner)
tokens = merge(winner, in: tokens)
let newToken = winner.left + winner.right
if tokenToId[newToken] == nil {
tokenToId[newToken] = idToToken.count
idToToken.append(newToken)
}
}
specialTokens = definitions.enumerated().map {
.init(name: $0.element.name, text: $0.element.text, id: idToToken.count + $0.offset)
}In this snippet, specialTokens corresponds to markers we can insert in the training text to signal transitions, such as the end of a text or, in the case of Shakespeare, speakers or splays. (More on this when we discuss training).
Now that we have a way to translate text into tokens, the next topic is creating Transformers, which we will cover in Part II. (There’s more than meets the eye).
N.B.: If you would like to see the code, it’s available here: https://github.com/indie-hub/shakegpt.git
