optimus prime
In previous post, i learnt what inference is. Before this, i thought when i ask ChatGPT to suggest me a movie, all it does is send the request to one of it's GPU, process it and send the response back. Boy was i wrong. There is some hardcore engineering involved here. But we'll persist. So today we go a little deeper.
We'll start with a simple prompt: "The cat sat". Our model should be able to predict "on the mat" and return "The cat sat on the mat" to us. As we will see there are multiple steps involved here.
1. Model
Model is a set of learned parameters. This parameters (or weights) can be as large as trillions. While model is trained, we repeatedly feed it text and ask it to predict the next word. Initially model might be terrible at predicting, but it keeps on learning by adjusting it's weights in the neural networks to make prediction more insightful. Eventually, this weights are the model.
Now we have a model which is a large set of numbers. But we still don't have a way to generate text.
2. Tokens
Model does not understand text or words. It needs to numbers to make predictions. So we need to convert the words in numbers and then convert numbers back to word for output. This is achieved through tokenization. At a very high level, tokens are like IDs associated with a word. For example
"The" → 101
"cat" → 102
"sat" → 103
"on" → 104
"the" → 105
"mat" → 106
But we can have multiple words in our vocabulary and it is not feasible to have IDs for each of them. User can also send any kind of gibberish like "thecatsatonthemat" and that would not be present in our vocabulary. So we break them down in subwords and then assign them an ID. This subwords can contain letters, whitespaces, punctuations etc. For example:
"inference"
can become
["in", "fere", "nce"]
Here each subword in the list becomes a token and then gets an ID assigned to it. Assume our vocabulary is:
Vocabulary
"the" → 123
" cat" → 456
" sat" → 789
"." → 42
"ing" → 931
"tion" → 817
then "The cat sat." becomes
[123, 456, 789, 42]
This IDs are what we feed into the model. It is important to understand that token IDs mean nothing. If "dog" = 101 and "cat" = 102, then it does not mean "dog" comes before "cat" or that "dog" is 1 unit away from "cat". They are just numbers we define so that we can process the text and make predictions using our learned weights. Tokens also determine context length of our model. If a model support 8192 tokens, it does not support 8192 words. It word can be broken down into multiple tokens as we have seen. So the model's context window is measured in tokens, not characters or words.
3. Embeddings
Now we have
"I made him an offer he can't refuse"
|
v
tokenizer
|
v
[10, 20, 30, 40, 50]
This numbers are not mathematic representations. They are merely identifiers for text in vocabulary. 30 doesn't mean "inference" is somehow three times more meaningful than 10. So we need to turn token IDs into vectors. That is done by embedding layer.
Imagine our vocabulary has 5 tokens:
0 <pad>
1 cat
2 dog
3 car
4 banana
The model maintains a embedding matrix. Say the embedding dimension is 4:
dimensions
1 2 3 4
|----------------------------|
cat │ 0.21 │-0.73 │ 0.42 │ 0.11 │
dog │ 0.19 │-0.68 │ 0.39 │ 0.14 │
car │-0.51 │ 0.22 │ 0.71 │ -0.32 │
banana │ 0.81 │ 0.14 │-0.23 │ 0.55 │
|____________________________|
So:
"cat"
↓
[0.21, -0.73, 0.42, 0.11]
and:
"dog"
↓
[0.19, -0.68, 0.39, 0.14]
These vectors are called embeddings. For each token, we get it's vector embeddings. and collectively for all the tokens we get a multidimensional matrix like:
[
[0.12, 0.51, -0.32, ...],
[0.21, -0.73, 0.42, ...],
[0.67, 0.11, 0.09, ...]
]
So using the embeddings, we got the vectors on which the neural networks can operate. This embeddings are processed by Transformer.
4. Transformers
The Transformer takes the vector representation of every token and repeatedly updates those representations based on the other tokens in the sequence. Suppose we have:
"The cat sat"
After embedding, We produce
The -> vector A
cat -> vector B
sat -> vector C
The Transformer processes them together:
[A, B, C] -> Transformer
and updates each one of them to produce new updated vectors. The vectors now contain information influenced by their context. A token needs token because on it's own, a token can mean multiple things. For example, "Bank is closed" and "Bank of river is steep". Here, "Bank" has very different meanings influenced by their context. Transformer allows representation of bank to incorporate information of it's surroundings. This is achieved through self-attention.
5. Self-attention
Self-attention lets tokens essentially ask: "Which other tokens should I pay attention to?" Suppose we have "The cat sat on the mat because it was tired.". Model needs to figure out "it" refers to here. Conceptually:
"The cat sat on the mat because it was tired."
^
|
"it"
|
|---------|---------|
| |
cat mat
0.7 0.1
So model calculates relationship between tokens. For every token, the model creates three vectors: Query, Key, Value. Usually abbreviated as Q, K, C. For every token we have:
Token vector
|----> Q
|----> K
|----> V
Query: "What information am I looking for?" Key: "What information do I contain / what am I relevant to?" Value: "Here's the actual information you can take from me." Model decides what's relevant by roughly doing Q * K and producing a similarity score.
A single transformer layer can have multiple multiple steps inside it. A typical block looks like:
Input
|
V
Self-Attention
|
V
Residual + Norm
|
V
Feed-Forward Network
|
V
Residual + Norm
|
V
Output
A LLM has multiple such transfers. This is why LLMs need expensive GPUs. They need to do expensive computations multiple times and GPUs are extremely good at performing huge numbers of parallel numerical operations.
6. Logits:
After all the calculations through the Transformers, model makes final prediction. It produces a score for every token in its vocabulary. Suppose our vocabulary has 50,000 tokens. The model produces something like:
[
1.2,
-0.4,
3.7,
0.1,
...
8.2,
...
]
That's a vector with 50,000 numbers. Those numbers are called logits. The higher the logit, roughly, the more strongly the model favors that token. However, Logits aren't probabilities. We convert them into probabilities using softmax. Suppose the model produces:
"The cat sat"
on 5.2
down 2.1
the 1.8
because 0.3
banana -2.1
After softmax, perhaps:
on 0.91
down 0.04
the 0.03
because 0.02
banana 0.00
Now we have a probability distribution. Which token we chose is decided by Sampling. Simplest strategy is to chose the one with highest probability. This is called greedy decoding. Using our above example, we chose "on". Now our sequence becomes "The cat sat on".
And we do the whole process again. Generation is essentially a loop. This is called autoregressive inference because the model uses its own previous outputs as inputs for future predictions.
Imagine generating 1000 tokens from 10,000 tokens prompt. This is a serious computational problem. This is solved by KV cache. I'll leave it for the next time.