Building your own Shakespearean GPT - a ChatGPT like model: Live with the R&D Team

A live tutorial demonstrating how to train a small, Shakespeare-style GPT (decoder-only transformer) in the Wolfram Language. The presenter covers tokenization of Shakespeare plays, building a vocabulary (~30k tokens from ~100k lines), token and positional embeddings, masked multi-head causal self-attention, residual connections, layer normalization, dropout, feed-forward layers, and assembling a transformer with 4 attention blocks, embedding size 128, and 4 heads. Training used cross-entropy loss with Adam, trained for 3 epochs on CPU with batch size 16 (10% validation split); training took ~1h20m. The talk also shows sampling strategies (greedy, temperature, top-k), discusses improvements (byte-pair / subword tokenization, larger models, GPUs), and compares outputs to OpenAI GPT-3/4 via an OpenAI Wolfram link.
Questions this video answers
About 100,000 play lines (~8 MB stored in a Wolfram Cloud object).
About 30,000 distinct tokens.
Embedding size 128 with 4 attention heads (key dim = 32).
Four attention blocks; trained for 3 epochs (rounds).
Batch size 16, validation split 10% (~10,000 lines), optimizer Adam, training on CPU took ~1 hour 20 minutes for 3 epochs.
Chapters
So, hello everyone. Thank you for joining me today. I'm gonna talk about a recent post that I published on the Wolfram community, so you can go there and download the notebook that I'm going to present today. So, recently Stephen Baufram wrote a blog post on what is ChatGPT doing and why does it work? Where he explains how ChatGPT and neural networks operate, and he discusses the limitations of GPT models when not assisted by external computational tools. So here we will use a different approach. We will walk through the technical details of training our own GPT model, even if it's going to be a small model using the Wolfram language. So let's get started. So first of all, what is a GPT? We hear a lot nowadays. ChatGPT even yesterday was announced the GPT 4. But many people don't know what these three letters mean. So GPT stands for Generative pre trained transformer, which is a type of deep neural network architecture based on a transformer that that was introduced in 2017 by Google in a paper called all you need is attention. And basically GPT is a machine learning model that generates text based on on the prompts it is given. Yeah, I mean there are like for example the GPT4 that was announced yesterday. It's a multimodal model that also can combine text, it can input text as well as images. But basically the architecture is going to be fairly the same. I see a first kind of request, if we can get this Shakespearean GPT model to use some temporal metaphors, maybe at the end we have some spare time and we can play as a disclaimer, this is more as a. It's going to be a pedagogical tutorial for training a GPT, a small GPT model. So we don't expect to have this GPT model to write an entire Shakespeare play from scratch. Although we will see some ways to use some more powerful GPT models like GPT3 or GPT4. So to give you a quick summary on how GPTs work, the generative part means that the can generate a sequence of tokens to be specific, while the pretty net part means that they already have been trained on tons of data from text, from books, from Internet and other sources. And finally, the transformer part refers to this type of neural network used in GPT, which specifically gives decoder origins from a neural network, which we will see in more details in the following parts. Some examples of large language models are the GPT models from OpenAI. As I mentioned yesterday, GPT4 was announced which is multimodal chatgpt as well codex, which runs behind the scenes. In the GitHub copilot there is Google's Lambda and many more different large language models. So these models are special because they are very large billions of parameters in terms of gigabytes of text or images in the case of the GPT4. So here we won't be able to train one of such large language models, but in our computer, on a personal computer we can train a small version of it having the same architecture. And for it for having some fun, we were going to try to train it using all the place from Shakespeare. So the first part is going to be to convert the input text from Shakespeare's place into a sequence of tokens. So when working with text data we often use a process called tokenization that's to break down the text into small pieces or tokens. And once we have these tokens we need a way to represent them numerically in our machine learning models. To do this we use a vocabulary that maps each token to a unique integer. And for this tokenization there are like different methods. Some more advanced techniques are the byte per encoding or the word piece. You can take them if you're interested, but their underlying principle is going to be the same as the one that we are going to use, which is the one word tokenization. So we will assign to each word to different token and also add the punctuation marks plus the spaces in our vocabulary. So for getting all the Shakespeare Place, I've already created a list of all the lines from Shakespeare Place and I put it into a Wolfram cloud object and so you can directly download it. It's only eight megabytes around eight megabytes. And here there are like some random sentences from short displays. And I guess that some of you have read some of the place and it feels quite familiar. So then here these are strings, we can select only the lines that we're interested in, which are the lines that the characters from the place are talking to each other. And for this we can use the select and stringstart skew to also get only the lines where they're like characters talking. And also then we can use the Stringtech tool to remove these quotation marks. And once we have these cleanup lines we can split them into one word tokens using string split with a word boundary. And another thing that we can do is to add this so called end of sentence token which will provide context later on and indicate the model when to stop a sequence. And avoiding generating Repetitive sequences. So that's why we are appending this end token at the end of each line. So we end up having around 100,000 play lines from Shakespeare. And these are like five random sequences of such lines already tokenized. As you can see there are some punctuation characters as well as the spaces. Now we can, with this list of tokens we can create the vocabulary which is going to be all the distinct words or punctuation characters that appear in the text. For this we can simply use union and we get around 30,000 such tokens, which is comparable to the current common English words used for example in the built in function wordlist, which is a bit more than 39,000. As a comparison we can also check what's the vocabulary size of one of the GPT models from OpenAI, for example GPT2 which was released in 2019. I think there is around 50,000 support tokens. They use support tokens in order to reduce the size of the vocabulary as well as being able to also to process non English words and other strange characters from Internet. So in the Wolfram Neural net repository we have already the GPT model available. So you can directly get it on the cloud or on a desktop notebook by using that model. And we can extract all the tokens there by getting all the labels from the output, for example and we can see that there are a bit more than 50,000 tokens. So we can also display some of these words from check displays. And we can see that nowadays these words are much less common, some of them at least. So the next part I'm going to explain is what is the tokens encoder and decoder? So in the Wolfram language we have the so called NET encoder that is going to determine an integer value for each token of the vocabulary. For example, the token love is assigned an integer value of 14,382. And if we want for example to create a vector that represents all the tokens from the vocabulary and we assigned a probability of zero to all of them except to the token love to one with unit vector we will with the native decoder we will get back the token. This type of kind of sampling, it's called greedy decoding since we are always taking the highest probability to make the final prediction. Later on we will see other sampling methods in order to get a better sequence generated sequences of tokens or better generated text. So at this point I'm gonna first go on a higher level to describe what is a decoder only transformer that's like the model architecture that GPT's models use. And it's the following. So we have the input sequences. In our case it's going to be text sequence of tokens. There is at first an initial token or text and positional embedding which then is going to fit in this blue block here, which is an attention block. And attention block consists in a multi head, causal or masked self attention. Then there are a residual connection here, a line normalization which we will into more details later on that sometimes nowadays also found before the multi head castles of attention, a feedforward layer, another residual connection, another layer normalization that as I mentioned sometimes it found before it. And then this attention block is repeated. There's a stack of them, several of them. In our case we're going to use four of them stacked together. But in large language models they can be kind of dozens of them, much larger. Then there's this output embedding layer and finally there is this linear layer that will project embeddings into a size of the boca. And finally the softmax will give us the probability distribution for the probability of the next token to be generated. So let's go into more detail in each of these parts of the network. So we will start with the initial tokens and positional embedding layers. So an embedding is a way of representing a token as a vector of numbers. So each token is assigned to a unique vector which is learned by the transformer during the training. These embeddings capture the semantic and syntactic relationships between tokens and text. The dimension of this vector will determine the depth of the network, which in our small model will be 128. But it can be much larger. Like for example one of the latest larger GPT3 models from OpenAI. The Da Vinci model is 12,288 this dimension and I guess that for GPT4 it's even larger that allow them to accommodate even longer sequences of text and process those. So then there's the positional embedding. That's in addition to the tokens. Embedding transformers have this. They need to to include this positional embedding in order to encode the position information of each token in a sentence. This is important because in the order of words in a sentence carries meaning in natural languages. As we know, in our case the positional embedding will be learned during training. But we could also use the sinusoidal position embedding maybe in the future in the original transformer model in Future blog post I will explain how to build your own sinusoidal positional embedding. But here for simplicity, we'll use only the learnable embedding positional embedding. And at the end, as it was found in the original paper, the results are practically the same. So now we can combine both embeddings and we obtained this net graph using the function layer. So we can see here, this is the tokens embedding from a sequence of n tokens. It will generate this kind of array of vectors for each token of 128 elements. And the other one, it's the position embedding that at the end they are added to other. So we can represent this initial embedding with array plot. For example the famous sentence of Shakespeare to be or not to be. And here each row it's representing one of the tokens in the sentence and it contains 128 elements of each. And this is before training it. That means that after training this will look a bit different. So now we go into the masked multi head self attention block layer. So that's this yellow part here in the diagram and this is we will break it down into different pieces. So first we're going to see what is the self attention mechanism. So the self attention mechanism involves generating three different matrices called query, key and value. And each of them contain a query key and value vectors for each of the the tokens in the sequence. In this case I'm just using this, I am a student example sequence. And then we to obtain these vectors we multiply the initial sequence embedding by different learned weight matrices for the so called query with matrices key and value. Once the query which they are like also like learned during training, once the query key and value vectors have been generated, the self attention mechanism computes a weight for each token in the sequence by taking the dot product between the query vector of that token and the key vector for every each other token in the sequence. So that's basically the representation of the dot product of the query and the key vectors for each of the sequence. And then there is also a scaling factor here which is the square root of the dimensionality of the vector. This is to stabilize it, to stabilize the training. So this product can grow sometimes. So in order to keep the number smaller they figure out that this will stabilize the training. And then there is a softmax function to the scalar product to obtain a set of weights that sum to one. And these weights are then used to compute the weighted sum of the value Vectors here and that which represent the importance of each token in the sequence. And we will obtain this attention value matrix. So another final component, not the final but not important component for GPT models is the so called masked attention. So the masked attention is used to prevent the model from looking at future positions in the input. So basically what we are doing is to inside this softmax we are adding a triangular mask of this sort where there is like minus infinities. So basically the sentence is like all models are wrong. This will be like this dot product here. And then we are adding this mask, it will after the softmax all these values are going to be zero. So the model basically cannot predict the next token by looking at the future tokens for doing training. That's why sometimes it is called causal self attention as well. And finally the part of the attention block that refers to the the multi head self attention is these attention heads. The concept of attention heads are it's a way of parallelizing self attention mechanism to allow the model to attend to different parts of the input sequence in a more fine gained way. So basically they separate in different instances of a self attention mechanism that operates in parallel with other attention heads in the model. Each attention head has its own set of query keys and value weight matrices. And this will allow to focus on different aspects of the input sequence. For example, by using multiple attention heads, the transformer model can learn to attend two different parts of the input sequence at the same time which can improve the quality of the model and focus. One head might focus on the syntax and another one might focus on the semantics of the sentence. After the attention head has been computed its own set of waves, they are concatenated together. This would be like the schematics of this multi head self attention, they are concatenated together and then pass it into a linear layer that would generate the final self attention mechanism representation with the same initial dimension that we had. So typically the embedding size is set up to be divisible by the number of heads. In our case the Embedding size is 128 and the number of heads that we will use is 4. So the dimension of the key vector will be 32. And yeah, there is a question from Bob. Does attention matrix scale like of the order of n squared? Yes, that's because. That's why we have the dot product. That's true. So another component of the tension block is the dropout layer. The dropout layer works by randomly setting a fraction of the activations in the previous layer to zero during Each training iteration it has the effect of reducing the capacity of the model and preventing it from relying too heavily on any one feature or set of features. So that prevents also the overfitting during training. We have a built in dropout layer in the Wolfram language, as we will see. And also another important part is the normalization layer, which I think that this type of normalization was introduced in 2015 for in deep learning models for image recognition. And since then it has been ubiquitous in any deep learning model. So the normal system layer is used to normalize the output of the previous layer, which helps to prevent the activations from becoming too large or too small. So this normalization is done by centering and scaling the activations to have zero mean and unit variance respectively. So a final other important component is these so called residual connections which involve adding the input to a layer to the output of the layer before passing it through the next layer of the network. This has the effect of creating a shortcut or a skip connection that allows the gradient to bypass to the layer and flow directly to the next layer. So this role of residual connection is to address the problem of so called vanishing gradients which can occur in deep neural networks when the gradient signal becomes too small and to effectively update the parameters of the model so the model will stop converting. So they're an important component. So now we can put all these previous layers that in order to create the masked causal self attention block using the built in functions of attention layer, dropout layer, normalization layer. And this is the function for this self attention block using function layer. And we can see here for example some of the explainer components. This is the scaling factor. The tension layer has the dot product, the multi head is set to be true and also includes a causal mask. And on here we can see the dropout layer. This represents the residual connection as we can see in the net graph of it. This will be like this connection going all the way directly to the next layer. And this is the tension mechanism that is for the dropout layer. And here there is the normalization layer. All right, so the next component that the next layer that these models use is the so called feed forward layer. So the feed forward layer is a fully connected neural network that consists of two linear layers with an activation function in between. In our case we're going to use this Gaussian error level linear unit element wise layer. You can check in the demonstration. There are several other different activation functions that are nowadays used in deep learning. So the first linear layer maps the input into A higher dimensional space, while the second maps the output back into the regional dimensionality. So the purpose of the feedforward layer is to introduce non linearity and allow the model to learn more complex features from the output of the self attention mechanism. This non linearity is achieved through the activation function which applies a thresholding function to the output of the first linear layer. And we can construct it in this manner here. And we can see here there is also a residual connection plus a dropout normalization layer. And in between the two linear layers there is the activation function. In our case is this Gaussian error linear loop. So finally we have all the components that we need for building up the whole decoder only transformer that GPT models use. And here I added some other hyperparameters. How many of these attention blocks we want in this in mode? During the experiment I was doing, I tried it with two and it seems that with four different blocks it also improved a bit the model. So we add together the different components that we saw, the embedding of the tokens and the positional embedding and then there is a stack of four different attention blocks. And finally there is the output that we add a linear layer and a softmax layer at the end to generate the distribution of predicted tokens. So now we only need to create the training model and that's by adding a cross entropy loss layer as the final layer of the network in order to calculate the loss between the predicted output and the true output using NET train. And basically the cross entropy loss is a standard loss function used in many natural language processing tasks such as language translation, text summarization. The loss is calculated by comparing the predicted output of the true output using a probability distribution. And we will see that the cross entropy is then back propagated through the network to update the weights of the network in order to minimize the loss. This is done using an optimization algorithm such as the stochastic gradient descent. We will use a similar one called Atom and in our case which iteratively updates the weights of the network to reduce the loss. So the final training net, it's gonna be like this here we can see that this is the decoder only transformer and this is the gloss entropy layer. So now we are in a position that we can train our network, our GPT model. We need to specify the model that we want to train and then include all the lines from checks we place. Here I specify that I only wanted to train it for three rounds to training on my laptop in a reasonable time. Then the position set was 10% of the whole lines. I think it's around like 10,000 lines for the position set. And I specified a batch size of 16 here. So in this case it took me one hour and 20 minutes to train these three rounds of training with all these lines. And you can see that I could have kept the mobile training since it was still like learning that the loss was going down as well as the validation loss. So there was no overfitting. So basically both the orange and the blue line, they should keep converging. But anyway we can already like try, try out the resulting net. I also save it into the Wolfram cloud so you can directly try the trainetnet without the need to recompute. Of course, if you had a GPU in your laptop, you can try to compute to train it in gpu. It's going to be much faster. I use the CPU on my own personal laptop. So this is the final render net. Now what we need to do is to create a predictor. So that's basically we're going to generate full sentences by iteratively getting the next token prediction from our model. At each iteration we will append the predictable back into the input and we will repeat this process in a auto regressive manner. So here I want to mention that this is not the most efficient way of creating the predictor. Since we are kind of computing over and over certain parts of the network during the generation of the sequence. We could, if a speed is an important thing for our application, you could use the net unfold built in function in order to make it more efficient during the decoding process and avoid computing over and over certain representations. So here we can use the. Since it's fast enough for kind of toy examples, we can use nest by list and tell the duplicate to stop whenever it encounters the end of sentence token or whether the sentence is longer than 20 tokens. So and I am initializing the sequence with this token, the O token so it generates oh my lord comma I am a man. A man is a man. So it starts. It looks like a bit repetitive. That's mainly because our model is a small GPT model and it's piroting. But also we can improve these results by not just picking the highest probability or token each time. This is called the degree sampling or degree research. But we can use some other most apparent sampling methods. So we can introduce some stochasticity or randomness during our generations by sampling from the probability distribution. Instead of being greedy, this will allow us to Generate different sentences giving the same input and also we can obtain better results. A bit more creative Here there's a link on a good tutorial explaining the different sampling methods. The most common ones are top k top probabilities and also the temperature hyperparameter. If we increase it, it will generate a bit more creative sentence in a model. So one can specify the sampling method directly on our model by indicating that we want for example to set up a temperature of 0.7 we will get such sequences oh my dear sweet fellow, exclamation mark. Or alternatively we can also pick up the top three probabilities and start from the same token we get what is the matter of my heart question mark. So now we can generate a batch of such sentences. But first here I'm randomly picking up certain words that start with uppercase in Shakespeare plays. Also keep in mind that if you introduce here a token that is not in shape, the model will not the token will not be part of the vocabulary and it will not be able to run. So this is a sample of 30 such lines. Most of them are still like nonsensical but they look more. Some of them look like more in. In the style of Shakespeare and there are likes even some of them are like a bit funny. So I will stop here improving my model. But here there are like certain improvements that we could do in order to make it more coherent to obtain a more coherent shapes deploy. Yes, there is a question from Bob saying that is there a better tokenizer? Yes, that this bipair encoding tokenizer that will generate soup board tokens will make a better kind of more suitable for having a more robust tokenizer. And that's the. And I think in the Wolfram language we have included support tokenization as well for the purpose of explaining of this tutorial. I thought it was more easily also for training it was. I think it was a bit faster using one word tokens instead. But yeah, definitely using a support tokenization would be the next step to make a better tokenizer. I think so one thing that we can do is to scale up the model so we increase the number of the dimension of the embeddings as we saw that we use it as more embedding depth dimension. It was 128 but in large GPT models they use several thousands. We could also increase the number of the decoder blocks, increase the number of heads in the causal multi heads of attention, increase the length of the training sequences. Putting together multiple lines of the model might be able to generate multiple lines that are coherent. But keep in mind that increasing all these parameters will also increase the training time of the model. So we should probably get access to some external GPUs. And finally, as a matter of comparison, I created an equivalent Shakespearean GPT directly using OpenAI GPT models. So here what I did is to use the OpenAI link which was created recently by Christopher Wolfram. So one needs to install this packlet and then also create a user account on the OpenAI side and generate your own API key. And then we only need to create a prompt that will do that. By default, the auto completion of this OpenAI link, it's going to use one of the most advanced GPT3 models called DaVinci model. And here this is going to be our problem. I'm asking GPT3 to write a fragment of a play that imitates Shakespeare's style. The characters are Romeo and Juliet and they talk about artificial intelligence and whether machines will be able to love one day. Here there is a fragment of a play from Shakespeare as a reference. So I'm giving. I'm copy pasting a fragment of one of Shakespeare's plays and in order to make use of the fact that these large language models are few short learners. So if you could give a specific task, it's better always to show them an example. And here they are the results. So they're not bad. Much better than our GPT model from scratch. And this is one that I tried that even give multiple lines. And then at the end I'm also adding some useful resources that if you are interested in learning more about GPT models or transformers in general, these are good resources to to keep learning. And this is the first one is it's a good place to learn machine learning using Wolfram language. This is a book from Etiennes Bernard. There is a nice chapter on deep learning methods using Wolfram language. The book is available for free online and you can also get a printed copy. There is this other blog post from Jalamar the Illustrated Transformer. It's a good place to also to understand better the transformer models. And then there is a couple of good reviews on the different attention mechanisms used nowadays in transformer models as well as a good review on different transformers like use it nowadays. It's a fast evolving field so it's hard to keep the base of all the different architectures that are being built. But if you understand once you understand one of the transformer architectures, it's fairly easy to understand the other ones. So then I'm explaining how to generate the header images at the end. Here I use it also the OpenAI link that packlet that we saw before and in this case I'm using the OpenAI create image instead of the auto completion one. And here I got some inspiration for the prompt using this prompt Hero app website and it will give you some tricks to get what you want. In this case I wanted a portrait of Shakespeare next to a robot that was a bit set and I got several different outputs and finally since yesterday OpenAI announced the GPT4 model on my ChatGPT plus account. I got early access on the GPT4 so I wondered try and compare with the same prompt as before to generate some output and to be honest it generated a much longer output and also quite. In line with Shakespeare style. Talking about kind of artificial intelligence and whether machines will be able to love one day. Yeah so if you have questions you can reply onto the community post. In the future I will share more kind of related topics on Transformers on the community and so thank you for joining me today and stay tuned for following discussions on Transformers. Thank you.
More from Guy

Cheryl Strauss Einhorn | The Human Edge: Smarter Decisions in the Age of AI | Talks at Google

Rachel Glennerster | Innovation and the Center for Global Development | Talks at Google

Rebecca Winthrop | Rethinking the Purpose of Education in the Age of AI | Talks at Google

Dr. Robert Wachter | A Giant Leap: How AI Is Transforming Healthcare... | Talks at Google






