Creating an LLM (little language model) from scratch — Part II
I remember being a kid and feeling fascinated by the Transformers cartoon. It was something fresh in that day and age. Thinking of autonomous robots that could transform themselves into cars and aeroplanes, and were here to help humanity (or kill, depending on the side). Wouldn’t that be cool in the future?, I thought. I guess my young self didn't expect humanity to be this close to that reality, as we are right now, thanks to recent developments in A.I.
In the previous post in this series, where I’m recounting how I created a little Shakespearean language model, I told you about the breakthrough moment of attention. The mechanism that unblocked the downpour of recent technological innovations. But attention is only half of the story. The complete story is something called a transformer. (This one, however, does not morph itself into a car)
So, what is a Transformer
A transformer is attention plus a feed-forward network with an activation function. And it’s all that it is. Now, if you are wondering what an activation function is, or even a forward network, I’m about to explain it.
If you read my previous post, you’ll remember that attention gives AI models relationship data. In my example of travelling from London to Edinburgh, having nothing but a topological map of Great Britain, attention provides contextual information, which is useful when picking a road to travel. Attention can tell you that Edinburgh is closer to Manchester than to Glastonbury, and that Glasgow is west of Edinburgh. But what should I do with this information? You probably know from attention that the Manchester bit is more important than the Glastonbury one, but what should I do with it? That’s what the feed-forward network with an activation function is all about.
Imagine that you are a feed-forward network. You reach a crossroad, and you consult your attention notes. After that, you take all notes you may have, and you expand with temporary notes (which may even be stupid, like saying that Glasgow is to the east of Edinburgh — we’ll see why later), and then you call on the assistance of the activation function. Magically, red and green lights appear on each road of the crossroad. Oh, that’s a huge help, you may say. With the traffic light, you now know that the green roads are good to go, while you should avoid the red ones. (Think of this step as having a personal Gandalf shouting “You shall not pass!”, on the roads that you should not travel.). You then keep on the map the notes about the green roads, and pick one to travel.
We call it feed-forward because it acts in layers (expand, activate, reduce), and each layer feeds the next (the values from one are used on the next). And, as magical as it may seem, there is nothing mystical about an activation function. It is a mathematical function that transforms values non-linearly. The most brutally simple version, ReLU, cuts everything below 0. Yes, that’s it. (There are, of course, more sophisticated functions, like the one used in my Shakespearean generator, which is based on a Gaussian, but the principle is the same: trim or smooth values).
And so you take your embeddings, run them through attention, expand, filter the values, reduce, and you have your Optimus Prime (or the embeddings that made the cut) ready to be sent to a decoder (more on this later).
Now, what about words? How can Shakespeare write? Do you still remember that what you want to predict is the next most statistically meaningful word? So, you take a sentence, you collect the relationship data from those words (from attention), you expand it to more relations (notes), you smooth those that do not matter, you reduce again to a smaller set, and you pick the winner (just not right now; there’s still a bit missing).
Simple, innit?
“Freedom is the right of all sentient beings.”
Excuse me, did you say embeddings?
Yes, embeddings. Another fundamental part of what makes A.I. what it is today. So, what are they?
Remember that we talked about tokens, and how words and parts of words are converted to numbers (do Transformers think in numbers?)? Well, those numbers are a map between bits of words. For example: 1 maps to ‘a’, 2 maps to ‘b’ and so on.
1 -> 'a'
2 -> 'b'
3 -> 'c'
...
256 -> 'th'
257 -> 'sa'
...Also, I alluded to the fact that attention creates relationships between data, which is the same as saying that attention creates relationships between tokens. But mapping one token (number) to a letter or part of a word doesn't give much information. We can compute the distance between two tokens, but it will not allow us to, for example, create clusters where we can say that cat is closer to tiger, and that red is closer to cardinal and far away from the blue band group. And all because we are working in a single-dimensional space (a line).
Get Bruno Oliveira’s stories in your inbox
Join Medium for free to get updates from this writer.
To enrich data, and let these clusters form, we use embeddings. So an embedding is a vector (an array) in a multidimensional space. Multi… Dimensional… Space.
Imagine that you want to map the natural numbers, 1, 2, 3, ..., to vectors in the N² space, (0, 0), (0, 1), (1, 0), … You can say that 1 maps to (0, 0), 2 maps to (0, 1), 3 maps to (1, 0) and so on (And although both spaces are infinite, they contain the same number of elements. Yes, infinity is weird, but that’s a discussion for another day.)
1 -> (0, 0)
2 -> (0, 1)
3 -> (1, 0)
...
256 -> (13, 42)
257 -> (13, 43)
...Now, the only difference to our purposes is that an embedding is not a bi-dimensional vector, but a multi-dimensional one (and real). Instead of mapping 1 to (0, 0), we associate it with (e₁, e₂, e₃, …, eₙ), where n is the size of the embedding. So, for an embedding of size 4, for example, you would have a map like
1 -> [ 0.42, -0.17, 0.83, 0.11]
2 -> [-0.66, 0.24, 0.05, 0.71]
3 -> [ 0.13, 0.91, -0.38, -0.22]
...
256 -> [-0.47, -0.54, 0.62, 0.29]
257 -> [ 0.78, -0.09, -0.31, 0.48]
...The embeddings in this example are random and meaningless, which, spoiler alert, is exactly how a model starts.
With a multi-dimensional vector, richer connections can surface, and the above clusters can form, all of this thanks to our good old mathematical operators like sums, subtractions, multiplications, divisions and functions such as distance or dot product (please don’t let the mathematics deter you). Embeddings are really a map to the stars.
Bringing all together
How does this all come together? Well, we start with a vocabulary, which is how we map words to tokens.
We then associate those tokens with embeddings (note that while the embedding can move in space, one token is always mapped to the same embedding).
After we have the embeddings, we feed them into a transformer.
The transformer takes the embeddings, gets the importance of these embeddings from attention (the relationships), expands the embeddings to a larger space (enriching them even more), suppresses or reshapes parts of them, and reduces them to their normal size.
How to go from the Milky Way to Andromeda? You take the map to the stars, you get the relationship between the stars, you expand the map (via potential connections to wormholes and multiverses), you cut those that do not matter, you shrink it back to a normal map, and the next light-year will be presented to you. (Almost… But, yes, only one light-year at a time nonetheless).
This is, in a nutshell, an LLM.
Below is a snippet of a Transformer as written in my Shakespearean generator (don’t mind the head count, or the dropout for now).
final class TransformerBlock: Module, UnaryLayer {
@ModuleInfo private var attention: MultiHeadSelfAttention
@ModuleInfo private var feedForward: FeedForward
private let norm1: LayerNorm
private let norm2: LayerNorm
private let dropout: Dropout
init(
embeddingSize: Int,
dropoutProbability: Float,
headCount: Int,
qkvBias: Bool
) {
self._attention.wrappedValue = MultiHeadSelfAttention(
embeddingSize: embeddingSize,
headCount: headCount,
dropoutProbability: dropoutProbability,
qkvBias: qkvBias
)
self._feedForward.wrappedValue = FeedForward(embeddingSize: embeddingSize)
norm1 = LayerNorm(dimensions: embeddingSize)
norm2 = LayerNorm(dimensions: embeddingSize)
dropout = Dropout(p: dropoutProbability)
}
func callAsFunction(_ input: MLXArray) -> MLXArray {
// Pre-normalise before attention, then preserve the original input
// through the first residual connection.
var x = norm1(input)
x = attention(x)
x = dropout(x)
x = x + input
// The second residual begins after attention so the feed-forward
// network adds its contribution to the already contextualised tokens.
let residual = x
x = norm2(x)
x = feedForward(x)
x = dropout(x)
x = x + residual
return x
}
}Small and simple. However, you know, the devil is in the details, and I left some of them dangling in the text. But I'll cover those in Part III, where we will take our little language model to the gym for some training.
N.B.: If you would like to see the code, it’s available here: https://github.com/indie-hub/shakegpt.git
