This explanation is adapted from the walkthrough by Brendan Bycroft, under the MIT license.
Introduction #
The demo uses nano-gpt, a small model with about 85,000 parameters.
Its goal is a simple one: take a sequence of six letters:
C B A B B C
and sort them in alphabetical order, i.e. to "ABBBCC".
We call each of these letters a token, and the set of the model's different tokens make up its vocabulary:
| token | A | B | C |
|---|---|---|---|
| index | 0 | 1 | 2 |
From this table, each token is assigned a number, its token index. And now we can enter this sequence of numbers into the model:
2 1 0 1 1 2
In the 3d view, each green cell represents a number being processed, and each blue cell is a weight.
-0.7 0.4 0.8 being processed -0.7 0.7 -0.1 weights
Each number in the sequence first gets turned into a 48 element vector (a size chosen for this particular model). This is called an embedding.
The embedding is then passed through the model, going through a series of layers, called transformers, before reaching the bottom.
So what's the output? A prediction of the next token in the sequence. So at the 6th entry, we get probabilities that the next token is going to be 'A', 'B', or 'C'.
In this case, the model is pretty sure it's going to be 'A'. Now, we can feed this prediction back into the top of the model, and repeat the entire process.
Preliminaries #
The visualization follows inference: running a trained model to produce a prediction.
This guide focuses on inference, not training, and as such is only a small part of the entire machine-learning process. In our case, the model's weights have been pre-trained, and we use the inference process to generate output. This runs directly in your browser.
The model showcased here is part of the GPT (generative pre-trained transformer) family, which can be described as a "context-based token predictor". OpenAI introduced this family in 2018, with notable members such as GPT-2, GPT-3, and GPT-3.5 Turbo, the latter being the foundation of the widely-used ChatGPT. It might also be related to GPT-4, but specific details remain unknown.
This guide was inspired by the minGPT GitHub project, a minimal GPT implementation in PyTorch created by Andrej Karpathy. His YouTube series Neural Networks: Zero to Hero and the minGPT project have been invaluable resources in the creation of this guide. The toy model featured here is based on one found within the minGPT project.
Embedding #
We saw previously how the tokens are mapped to a sequence of integers using a simple lookup table. These integers, the token indices, are the first and only time we see integers in the model. From here on out, we're using floats (decimal numbers).
Let's take a look at how the 4th token (index 3) is used to generate the 4th column vector of our input embedding.
We use the token index (in this case B = 1) to select the 2nd column of the token embedding matrix on the left. Note we're using 0-based indexing here, so the first column is at index 0.
This produces a column vector of size C = 48, which we describe as the token embedding.
And since we're looking at our token B in the 4th position (t = 3), we'll take the 4th column of the position embedding matrix.
This also produces a column vector of size C = 48, which we describe as the position embedding.
Note that both of these position and token embeddings are learned during training (indicated by their blue color).
Now that we have these two column vectors, we simply add them together to produce another column vector of size C = 48.
We now run this same process for all of the tokens in the input sequence, creating a set of vectors which incorporate both the token values and their positions.
Feel free to hover over individual cells on the input embedding matrix to see the computations and their sources.
We see that running this process for all the tokens in the input sequence produces a matrix of size T x C. The T stands for time, i.e., you can think of tokens later in the sequence as later in time. The C stands for channel, but is also referred to as "feature" or "dimension" or "embedding size". This length, C, is one of the several "hyperparameters" of the model, and is chosen by the designer to in a tradeoff between model size and performance.
This matrix, which we'll refer to as the input embedding is now ready to be passed down through the model. This collection of T columns each of length C will become a familiar sight throughout this guide.
Layer Norm #
The input embedding matrix from the previous section is the input to our first Transformer block.
The first step in the Transformer block is to apply layer normalization to this matrix. This is an operation that normalizes the values in each column of the matrix separately.
Normalization is an important step in the training of deep neural networks, and it helps improve the stability of the model during training.
We can regard each column separately, so let's focus on the 4th column (t = 3) for now.
The goal is to make the average value in the column equal to 0 and the standard deviation equal to 1. To do this, we find both of these quantities (mean (μ) & std dev (σ)) for the column and then subtract the average and divide by the standard deviation.
The notation we use here is E[x] for the average and Var[x] for the variance (of the column of length C). The variance is simply the standard deviation squared. The epsilon term (ε = 1×10-5) is there to prevent division by zero.
We compute and store these values in our aggregation layer since we're applying them to all values in the column.
Finally, once we have the normalized values, we multiply each element in the column by a learned weight (γ) and then add a bias (β) value, resulting in our normalized values.
We run this normalization operation on each column of the input embedding matrix, and the result is the normalized input embedding, which is ready to be passed into the Self-Attention layer.
Self Attention #
The self-attention layer is perhaps the heart of the Transformer and of GPT. It's the phase where the columns in our input embedding matrix "talk" to each other. Up until now, and in all other phases, the columns can be regarded independently.
The self-attention layer is made up of several heads, and we'll focus on one of them for now.
The first step is to produce three vectors for each of the T columns from the normalized input embedding matrix. These vectors are the Q, K, and V vectors:
- Q: Query vector
- K: Key vector
- V: Value vector
To produce one of these vectors, we perform a matrix-vector multiplication with a bias added. Each output cell is some linear combination of the input vector. E.g. for the Q vectors, this is done with a dot product between a row of the Q-weight matrix and a column of the input matrix.
The dot product operation, which we'll see a lot of, is quite simple: We pair each element from the first vector with the corresponding element from the second vector, multiply the pairs together and then add the results up.
This is a general and simple way of ensuring each output element can be influenced by all the elements in the input vector (where that influence is determined by the weights). Hence its frequent appearance in neural networks.
We repeat this operation for each output cell in the Q, K, V vectors:
What do we do with our Q (query), K (key), and V (value) vectors? The naming gives us a hint: "key" and "value" are reminiscent of a dictionary in software, with keys mapping to values. Then "query" is what we use to look up the value.
Software analogy Lookup table: table = { "key0": "value0", "key1": "value1", ... } Query Process: table["key1"] => "value1"
In the case of self-attention, instead of returning a single entry, we return some weighted combination of the entries. To find that weighting, we take a dot product between a Q vector and each of the K vectors. We normalize that weighting, before finally using it to multiply with the corresponding V vector, and then adding them all up.
Self Attention Lookup table: K: V: Query Process: Q: w0 = . w1 = . w2 = . [w0n, w1n, w2n] = normalization ([w0, w1, w2]) result = w0n _ + w1n _ + w2n *
For a more concrete example, let's look at the 6th column (t = 5), from which we will query from:
The {K, V} entries of our lookup are the 6 columns in the past, and the Q value is the current time.
We first calculate the dot product between the Q vector of the current column (t = 5) and the K vectors of each of the those previous columns. These are then stored in the corresponding row (t = 5) of the attention matrix.
These dot products are a way of measuring the similarity between the two vectors. If they're very similar, the dot product will be large. If they're very different, the dot product will be small or negative.
The idea of only using the query against past keys makes this causal self-attention. That is, tokens can't "see into the future".
Another element is that after we take the dot product, we divide by sqrt(A), where A is the length of the Q/K/V vectors. This scaling is done to prevent large values from dominating the normalization (softmax) in the next step.
We'll mostly skip over the softmax operation (described later); suffice it to say, each row is normalized to sum to 1.
Finally, we can produce the output vector for our column (t = 5). We look at the (t = 5) row of the normalized self-attention matrix and for each element, multiply the corresponding V vector of the other columns element-wise.
Then we can add these up to produce the output vector. Thus, the output vector will be dominated by V vectors from columns that have high scores.
Now we know the process, let's run it for all the columns.
And that's the process for a head of the self-attention layer. So the main goal of self-attention is that each column wants to find relevant information from other columns and extract their values, and does so by comparing its query vector to the keys of those other columns. With the added restriction that it can only look in the past.
Projection #
After the self-attention process, we have outputs from each of the heads. These outputs are the appropriately mixed V vectors, influenced by the Q and K vectors.
To combine the output vectors from each head, we simply stack them on top of each other. So, for time t = 4, we go from 3 vectors of length A = 16 to 1 vector of length C = 48.
It's worth noting that in GPT, the length of the vectors within a head (A = 16) is equal to C / num_heads. This ensures that when we stack them back together, we get the original length, C.
From here, we perform the projection to get the output of the layer. This is a simple matrix-vector multiplication on a per-column basis, with a bias added.
Now we have the output of the self-attention layer. Instead of passing this output directly to the next phase, we add it element-wise to the input embedding. This process, denoted by the green vertical arrow, is called the residual connection or residual pathway.
Like layer normalization, the residual pathway is important for enabling effective learning in deep neural networks.
Now with the result of self-attention in hand, we can pass it onto the next section of the transformer: the feed-forward network.
MLP #
The next half of the transformer block, after the self-attention, is the MLP (multi-layer perceptron). A bit of a mouthful, but here it's a simple neural network with two layers.
Like with self-attention, we perform a layer normalization before the vectors enter the MLP.
In the MLP, we put each of our C = 48 length column vectors (independently) through:
-
A linear transformation with a bias added, to a vector of length 4 * C.
-
A GELU activation function (element-wise)
-
A linear transformation with a bias added, back to a vector of length C
Let's track one of those vectors:
We first run through the matrix-vector multiplication with bias added, expanding the vector to length 4 * C. (Note that the output matrix is transposed here. This is purely for vizualization purposes.)
Next, we apply the GELU activation function to each element of the vector. This is a key part of any neural network, where we introduce some non-linearity into the model. The specific function used, GELU, looks a lot like a ReLU function (computed as max(0, x)), but it has a smooth curve rather than a sharp corner.
We then project the vector back down to length C with another matrix-vector multiplication with bias added.
Like in the self-attention + projection section, we add the result of the MLP to its input, element-wise.
We can now repeat this process for all of the columns in the input.
And that's the MLP completed. We now have the output of the transformer block, which is ready to be passed to the next block.
Transformer #
And that's a complete transformer block!
These form the bulk of any GPT model and are repeated a number of times, with the output of one block feeding into the next, continuing the residual pathway.
As is common in deep learning, it's hard to say exactly what each of these layers is doing, but we have some general ideas: the earlier layers tend to focus on learning lower-level features and patterns, while the later layers learn to recognize and understand higher-level abstractions and relationships. In the context of natural language processing, the lower layers might learn grammar, syntax, and simple word associations, while the higher layers might capture more complex semantic relationships, discourse structures, and context-dependent meaning.
Softmax #
The softmax operation is used as part of self-attention, as seen in the previous section, and it will also appear at the very end of the model.
Its goal is to take a vector and normalize its values so that they sum to 1.0. However, it's not as simple as dividing by the sum. Instead, each input value is first exponentiated.
a = exp(x_1)
This has the effect of making all values positive. Once we have a vector of our exponentiated values, we can then divide each value by the sum of all the values. This will ensure that the sum of the values is 1.0. Since all the exponentiated values are positive, we know that the resulting values will be between 0.0 and 1.0, which provides a probability distribution over the original values.
That's it for softmax: simply exponentiate the values and then divide by the sum.
However, there's a slight complication. If any of the input values are quite large, then the exponentiated values will be very large. We'll end up dividing a large number by a very large number, and this can cause issues with floating-point arithmetic.
One useful property of the softmax operation is that if we add a constant to all the input values, the result will be the same. So we can find the largest value in the input vector and subtract it from all the values. This ensures that the largest value is 0.0, and the softmax remains numerically stable.
Let's take a look at the softmax operation in the context of the self-attention layer. Our input vector for each softmax operation is a row of the self-attention matrix (but only up to the diagonal).
Like with layer normalization, we have an intermediate step where we store some aggregation values to keep the process efficient.
For each row, we store the max value in the row and the sum of the shifted & exponentiated values. Then, to produce the corresponding output row, we can perform a small set of operations: subtract the max, exponentiate, and divide by the sum.
What's with the name "softmax"? The "hard" version of this operation, called argmax, simply finds the maximum value, sets it to 1.0, and assigns 0.0 to all other values. In contrast, the softmax operation serves as a "softer" version of that. Due to the exponentiation involved in softmax, the largest value is emphasized and pushed towards 1.0, while still maintaining a probability distribution over all input values. This allows for a more nuanced representation that captures not only the most likely option but also the relative likelihood of other options.
Output #
Finally, we come to the end of the model. The output of the final transformer block is passed through a layer normalization, and then we use a linear transformation (matrix multiplication), this time without a bias.
This final transformation takes each of our column vectors from length C to length nvocab. Hence, it's effectively producing a score for each word in the vocabulary for each of our columns. These scores have a special name: logits.
The name "logits" comes from "log-odds," i.e., the logarithm of the odds of each token. "Log" is used because the softmax we apply next does an exponentiation to convert to "odds" or probabilities.
To convert these scores into nice probabilities, we pass them through a softmax operation. Now, for each column, we have a probability the model assigns to each word in the vocabulary.
In this particular model, it has effectively learned all the answers to the question of how to sort three letters, so the probabilities are heavily weighted toward the correct answer.
When we're stepping the model through time, we use the last column's probabilities to determine the next token to add to the sequence. For example, if we've supplied six tokens into the model, we'll use the output probabilities of the 6th column.
This column's output is a series of probabilities, and we actually have to pick one of them to use as the next in the sequence. We do this by "sampling from the distribution." That is, we randomly choose a token, weighted by its probability. For example, a token with a probability of 0.9 will be chosen 90% of the time.
There are other options here, however, such as always choosing the token with the highest probability.
We can also control the "smoothness" of the distribution by using a temperature parameter. A higher temperature will make the distribution more uniform, and a lower temperature will make it more concentrated on the highest probability tokens.
We do this by dividing the logits (the output of the linear transformation) by the temperature before applying the softmax. Since the exponentiation in the softmax has a large effect on larger numbers, making them all closer together will reduce this effect.