Predicting the next word in a sequence is the task of a language model. A language model is a form of self-supervised learning. Unlike most deep learning tasks, language models don’t have the targets of the problem provided to them through manual labelling. Instead, the data themselves contain the required labels. In our case, these labels are simply the next word in the sentence. There are many reasons you might want to undertake such a self-supervised learning problem, beyond the obvious ones. For instance, fitting a language model and then using transfer learning improves the performance of text classifiers enormously: they can understand the detail of the corpus they are predicting. This approach is known as Universal Language Modelling Fine-Tuning(ULMFit) and was pioneered by Jeremy Howard and Sylvain Glugger.
Whilst transformers are all the rage of the deep learning community at the moment, it’s difficult to find resources on language models in Pytorch and fastai. It’s even more difficult to gain an understanding of the entire language pipeline, from importing text data and tokenising to writing complex models. Thus, our task will be predicting future words in the Wikitext Tiny dataset, courtesy of Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher (2016, Pointer Sentinel Mixture Models). This dataset is available under a Creative Commons Attribution-ShareAlike License.
. It is a small sample of the corpus of over 100 million words collected from Wikipedia articles. This article combines the foundations of natural language processing, including tokenisation, along with state-of-the-art techniques for getting RNNs to work really well. This simultaneously addresses the topics covered in both Chapters 12 and 14 of Jeremy Howard’s fastbook, and elaborates on concepts and code that require further generalisation.
Data preprocessing
You might wonder why we’re using a smaller model. If we have the whole Wikitext data available (through fastai), then why not use that? There are two main reasons: whenever you begin a modelling task, always start with a subset of data. It allows you to set up your pipeline, iterate more quickly, and experiment more easily. Once you’ve got the whole thing working, then it’s a simple matter of substituting in the larger dataset, and waiting for the model training to finish.
We’ll start by reviewing the two main steps in converting text to something that can be fed into a model: tokenisation and numericalisation.
Tokenisation
Tokenisation is the process of converting a list of words into individual tokens. However, we can’t just break up the words on spaces. If you think about it, there are many quirks to the English language, and many strange symbols, that we need to account for. For instance, fastai uses a set of rules for tokenisation called spaCy. We can access this default tokeniser by using WordTokenizer.
We use the function coll_repr(collection, n) to display the first n items of this collection. We also have to wrap txt in a list, as the tokeniser takes a collection of documents, which we can express as a list. However, we can’t just leave it at this. As those of you familiar with language models will be aware of, we need to add special characters indicating the start of a new text: xxbos, where BOS stands for "beginning of stream". There are lots of other rules that also make it easier for a language model to learn, such as reducing the size of the embedding matrix by using special tokens for repeated characters, rather than simply repeating the token several times. This is handled by fastai’s Tokenizer():
Other special tokens include:
xxmaj: the next word begins with a capitalxxunk: an unknown token
Numericalisation
Once we’ve tokenised our input stream, we still need a way to categorise these tokens, much like we might embed categorical variables in an embedding matrix for tabular data. In fact, the process is almost identical: make a list of all possible tokens, representing the categorical variable (this is just our vocabulary) and then replace each token with its index or place in this list.
To do so, we actually need our tokenised input stream first, so we call our tokeniser from above: toks=tkn(txt). We then need to setup Numericalize: this is done using a class method, passing in our tokens as an argument. In this way, we construct the vocabulary of the corpus. After it’s been setup, we can then pass any input stream into it, to convert tokens into integers. We specify min_freq=0 because the default is 3; since our text is not even a sentence, this would mean that any word appearing less than three times would be replaced by xxunk.
If we were to pass in our text from above into num, what would we expect? Well, the whole point of the process is to convert tokens into integers, representing their index in the vocab. Let’s double check we get this.
You might be wondering why we have so many 0s. This is because our input stream is so small, so most of the tokens will be unknown. We need lots of batches to actually construct a large enough vocabulary. I recommend looking at Chapter 10 of the fastai book for an example of this.
Why would we want unk tokens in the first place? Surely it’s good to feed as much information to our model as possible? However, avoiding an overly large embedding matrix is useful, because it can use too much memory and slow down training. Additionally, the words that aren’t well represented in the corpus likely don’t have enough data to form representations on anyway.
Creating the dataloaders
Initially, we’re going to create the dataloaders for our language model essentially from scratch, and then later we’re just going to shove all our text into a Pandas dataframe and let fastai take care of it for us. This initial step is important for understanding how to feed data to a language model.
We extract the data in the usual way, and then use the built-in function L() to read the text data. This is essentially just a fancy version of a Python list, with a few advanced indexing techniques inspired by NumPy. Don’t worry too much about the semantics of this. You can recycle this code for different text files or CSV files in future. This simply concatenates both the training and test sets into one contiguous stream. This will enable us to split it up into significant and sensible batch sizes later.
We then can have a look at our text.
We’ll tokenise by splitting the text at every space " ".
Creating a list of unique tokens gives us the vocabulary of the corpus. Again, we use the built-in function L(). We can then map each token to its index in the vocabulary by enumerating over vocab. As we would expect, feeding our tokens in one at a time gives as a list of integers.
To use our language model, we need to make certain that the batches are organised in a specific way. In the future, we can use built-in fastai functionality to do this for us, but for now we’ll do it by hand.
The dataset is split into bs equally sized groups, representing bl mini-streams of sequences. Each batch has bs rows (batch size) and bl columns (sequence length). The first row of the first batch contains the beginning of the first mini-stream. The second row of the first batch contains the beginning of the second mini-stream. The first row of the second batch contains the second chunk of the first mini-stream, following on from where the last batch finished off. The model will "see" the text flowing from batch to batch.
Our group_chunks function here comes directly from fastai, and all it does is this reindexing. We do that by breaking up the text stream into m = len(dset) // bs batches of equal size, where m is the length of each batch.
You might wonder about what happens if the length of the text stream is not exactly divisible by the batch size. In this case, we can drop the last incorrectly shaped batch (the remainder when dividing by batch size) when creating the DataLoaders with drop_last=True. You might also question how the texts are going to be in continuous streams if the dataloaders is built on the concept of random provision of batches to the model. Because of this, we also pass shuffle=False.
Finally, we need to take the text streams and break them up into corresponding inputs and targets. To do this, we take the sequence length of our input, say 16 words, and index our words tokens from our current word up to 16 words in the future. Then, the target becomes the 16 words starting after the current word (offset by one). Make sure you understand the range of the loop in seqs below. We then specify a cutoff for the train/test split, being 80% of our data. Finally, we use DataLoaders as usual, except this time passing in our train and valid datasets directly using the method from_dsets, and remembering to pass in our batch size.
We can confirm that our targets are simply our inputs, offset by one element, by looking at the first item in the seqs list.
This is exactly what we want: we feed in a certain stream of words, and the self-supervised language model is tasked with predicting the next word in the stream.
Baseline RNN
Our initial model will still be a recurrent neural network, but it won’t have the interior cell structure of an LSTM. Instead, we’ll use a vanilla RNN, with only one layer. Our task will be feeding in the last n words at each time step, outputting a prediction, and using the hidden cell state to make a prediction when the next word is fed in. We loop over this process for as many words as we have in our sequence, minus one.
We use three layers in our RNN: an embedding layer, and two linear layers, the last linear layer outputting the predicted word. The only thing that is slightly unfamiliar here is the embedding layer. All it does is take a list of the indices, representing our tokens, and output the word embeddings. We also take care to initialise our hidden state, represented by h, to 0 in the initialisation. If we did this in the forward method, we would be throwing away the hidden state of the model after each sequence is passed in. This would defeat the whole purpose – the model wouldn’t have any clue where in the sequence it’s up to.
In our forward method, we initialise a list of outputs. We then loop over the sequence length, at each iteration concatenating the hidden state and the embedding layer applied to the current input (up to the current time-step along the second axis). We then apply the middle linear layer and a ReLU nonlinearity, and finally append the output of the final linear layer to outputs. We stack these outputs along the first dimension when returning them.
The only line that doesn’t quite make sense is the self.h=self.h.detach() line. This is necessary because of how RNNs work. By initialising the hidden state to 0 only at the start of a sequence (in our case, a sequence is an article), then the model may have to loop over itself several thousand times, for the thousands of words in the article. (Remember, this is because we advance one word at a time.) This essentially equates to a several-thousand-layer neural network, if you think about the unrolled representation of an RNN. If you haven’t spotted the problem with that yet, here it is: if we’re on the 8000th word, and thus the 8000th layer of the RNN, backpropagation requires us to calculate the derivatives all the way back to the first layer. As you can imagine, this takes a long time and a lot of memory.
Hence, we simply remove all the gradients except the gradients from our sl length input. If our sequence length is 16 or so, this makes it much more manageable. We do this with the detach method, which gets called after every loop of sl iterations. The calculation of the derivatives backwards through recurrent layers is called backpropagation through time (BPTT). The detach method is often referred to as truncated BPTT.
Notice also that we add a reset method, which we can access with a callback in creating our learner. All this does is automatically initialise h to zeros at the start of each epoch, and before validation. The reasoning for this is that we don’t want the model to have some hidden state from a golf article when it reads and then predicts the first word of an article on cooking.
We can then create our learner and call fit_one_cycle. Remember, language models take a while to train, so even one epoch with our simple model will take around 20 minutes. Note how we pass ModelResetter to cbs, standing for callbacks.
15% accuracy. Not bad, but let’s see if we can improve on that.
Creating the LSTM model
Whilst LSTMs can partly solve the issue of exploding or vanishing gradients in comparison to RNNs, there are still three key things we can do to ensure our LSTM is regularised and doesn’t overfit. These were all introduced in a seminal paper by Merity, Keskar and Socher. In order of importance, they are:
- Dropout
- Activation and temporal regularisation
- Weight tying
We’ll see what contributions the above techniques have compared to the performance of the vanilla RNN.
Dropout
We use dropout to regularise the neural network and prevent overfitting. It works by randomly zeroing out activations in the specified layer of the neural network, with probability p. This prevents over-reliance of the network on certain patterns of activation, allowing the neurons to cooperate better. The added noise improves the robustness of the model. This idea was actually introduced in an online lecture by Geoffrey Hinton, one of the founders of deep learning.
Hinton’s inspiration for dropout was a trip to the bank. "I went to my bank. The tellers kept changing and I asked one of them why. He said he didn’t know but they got moved around a lot. I figured it must be because it would require cooperation between employees to successfully defraud the bank. This made me realize that randomly removing a different subset of neurons on each example would prevent conspiracies and thus reduce overfitting."
However, we can’t just randomly drop out some activations and leave it at that. Think about the consequences for the next layer. Say we have 10 neurons originally, and throughout dropout we are left with 5. Then, passing these activations onto the next layer, we can imagine the next layer being rather confused. Previously, it was getting the sum of 10 positive (since we apply ReLU between layers) activations, and now it’s only getting half of that signal. That is, the scale is different.
If we want the scale to remain the same, we need to rescale. If p=0.5, for instance, then we expect half of the neurons to be zeroed out, which will approximately halve the information being passed to the next layer. We have about 1-p neurons left over. Thus, we divide by 1-p=0.5 to increase the size of these activations to the same level as before (as we are dividing by a number less than 1).
Activation and temporal regularisation
Weight decay is used in deep learning as a form of regularisation to prevent overfitting. Weight decay reduces the size of the model’s weights by adding a penalty to the loss, meaning that weights can’t become overly important compared to all the others. Similarly, there are two other forms of regularisation we can apply to recurrent neural networks: activation regularisation and temporal activation regularisation.
Activation regularisation (AR) aims to make the final activations of the model as small as possible. To do so, we simply square the activations and add them to the loss, multiplied by a hyperparameter alpha. Reducing the size of the activations achieves sparsity, which is desirable when we have large hidden layers that are inclined to overfit the data. Introducing such a sparsity penalty allows the model to generalise better. To make AR work, we apply it on the dropped-out activations, so that we don’t penalise activations zeroed out in regular dropout.
Temporal activation regularisation (TAR) works by ensuring that activations between consecutive tokens in the sequence are close together. This is because the activations shouldn’t change drastically after seeing just one new word in the sentence. Thus, we penalise the loss by adding the squared difference between consecutive activations. In contrast to AR, TAR is applied on the non-dropped-out activations, as we want the difference of actiativations between timesteps, which could be significantly reduced by dropout at training.
Weight tying
The AWD LSTM paper also introduced a technique known as weight tying. This relates to the symmetry of mappings in the model. In our language model, embedding the inputs is a mapping process that takes words to activations. Conversely, the final output layer takes activations to English words. Based on intuition, these mappings could be symmetric. We can enforce Pytorch to assign the same weight matrix to the embedding input and output layers as follows:
self.h_o.weight = self.i_h.weightAWD-LSTM
Let’s include all these tricks and see how much better our accuracy is. Recall, we are still using the same LSTM model architecture; we’re just doing little bits and pieces in between steps that make it an AWD-LSTM. Note the similarity between how we would normally construct an LSTM in Pytorch, and our model here. The key step is creating the layer nn.LSTM, where we take n_hidden inputs and output n_hidden outputs, with a certain number of layers represented by n_layers. Pytorch handles the rest for us. We also implement dropout here after the LSTM layers, and notice the line that instantiates weight tying.
Note that we return three things for each forward pass through the network: the regular output, the activations from the RNN layer before dropout, and the activations after dropout. This allows us to use the callback RNNRegularizer.
To use activation regularisation and temporal regularisation, we don’t actually pass it into the model. This is largely due to the fiddly nature of implementing it ourselves, and because it can be tacked on directly to the loss. Instead, we pass RNNRegularizer to our callbacks list when creating our learner, with the hyperparameters of alpha and beta for activation and temporal regularisation respectively.
You might be wondering why we created two learners. The first was created in the usual way: pass in the dataloaders, model, loss function, metrics and callbacks, and we’re done. The second is a simpler way to achieve this. We can let the TextLearner take care of the callbacks like ModelResetter and RNNRegularizer, which are done automatically.
After one epoch of training, our model is already substantially better than the multilayer RNN. If you have time to train for longer, I highly recommend doing multiple epochs, as the difference between the models will be even more apparent after doing so.
Training using fastai
Batches
Whilst having this knowledge of how tokenisation and numericalisation works in language models is important for debugging, we can actually use fastai’s inbuilt modules to do it for us. For our purposes, we don’t even need to pass a DataLoader object to DataLoaders. We can create a dataloaders using a Datasets object, which I find to be easier when dealing with Pandas DataFrames. First, we’ll concatenate our training and test texts into one big contiguous stream, as we did above.
Then, we’ll create the Datasets object. We will use this to create our dataloaders. Let’s break this code down line by line.
Obviously, splits are used to create the indices of our train and validation sets, where we use the range_of function to do so (hint: look carefully at where the brackets are placed). Then, we need to pass certain transforms to the Datasets class initialiser, with tokenisation and numericalisation among them. If you read the docs for Tokenizer, you’ll see that the 0 we pass in here is simply the name of the columns of text. (Print out df_all and you’ll see that is indeed the column name.)
However, there is also a strange operator from the Python standard library called attrgetter("text"). This essentially constructs a function to get a certain attribute from a class. For example, if boy is an instance of a certain class, and we define func = attrgetter("age"), then calling func(boy) will return boy.age. Thus, in our call here, attrgetter("text") will take the "text" attribute created in the tokenisation of our txt_cols and numericalise it.
Finally, we pass in all these arguments to our Datasets constructor. Notice that we specify the dataloader type as a language model, which means that our dataloader will perform the actions we have outlined above when we pass in additional arguments during the call to the class method Datasets.dataloaders(), as we do here:
What we do here is concatenate all the text samples (the rows of the dataframe) into one stream, split it into contiguous sequences of size bs (representing batch-size) and then go through each of these seq_len tokens at a time. And then we’re done! We can show a few batches in the usual way after we’ve created our dataloaders with dls.show_batch(max_n=3). Notice that the dependent variable on the right is simply the input stream offset by one token.
Language-model learner
We can save ourselves a lot of hassle by using a built-in language model learner in fastai that also uses an optimised AWD-LSTM architecture. It will also be a much bigger model than the one we used above. We construct the learner in pretty much the same way, but pass in AWD-LSTM as the model. We also include perplexity as a metric, which is simply the exponential of the loss, and is commonly used in language models.
We call fit_one_cycle on the learner, and once an epoch is complete, we’ll save our model to our local machine, so we can load it back later without worrying about going through all the above steps again. Warning: without a GPU, even just this one epoch will take a while. By prepared to wait for a couple of hours.
Amazing! Our accuracy has doubled over our own AWD-LSTM, and is nearly triple the accuracy of the vanilla RNN.
Saving and loading models
Call learn.save("final_model"). This will create a file in learn.path/models/ named "final_model". We can reload the weights of this model to the learner with learn.load("final_model").
Text generation
One of the reasons for using the language_model_learner rather than an ordinary Learner above was that language_model_learners have the .predict method for generating new text. To do so, we only need to specify the starting text, the number of words, and we’ll loop over this for the number of desired sentences.
All we need to do now is print them out.
print("n".join(preds))Michael Jordan was a professional golfer , who won a Golden Globe Award and the WWE Emmy Awards . Jordan was named to the NBA All - Star Game Team for that yearMichael Jordan was a professional basketball player . Jordan said that " i have a lot of support for my career " , and that he had changed his mind after being recalled to Utah . Jordan said he would not beThat, at least to me, is amazing. Our language model is obviously incorrect about the facts, but if you weren’t intimately aware about Michael Jordan’s biography, this automatically generated text could almost pass as a Wikipedia article. Using a relatively high value for Temperature increases the randomness of the model, and we can also pass arguments like min_p to not consider words below a certain threshold probability.
The predictions also have almost perfect grammar (as Jeremy Howard points out, the "i"s aren’t capitalised simply because it’s the tokenisation rules not to capitalise a single character). Our language model, in only a few hours of training on a CPU, has picked up on very core concepts that form the bedrock of the English language.
Conclusion
In this article, we covered how to build a language model with Pytorch and fastai, which is a model that can predict the next word in a sentence. We trained our model on several thousand Wikipedia articles, and improved accuracy by nearly triple, and used the LSTM to generate an article about Michael Jordan.
If you want to find out more about how to handle the Wikitext dataset (and NLP datasets in general), look at the brief fastai walkthrough here. A future direction for exploration would be taking the language model created in this article, and using it to train a language classifier model, much like in ULMFit.
References
[1] S. Merity, N. S. Keskar and R. Socher, Regularizing and Optimizing LSTM Language Models (2017), CoRR
[2] J. Howard and S. Ruder, Universal Language Model Fine-tuning for Text Classification (2018), CoRR
[3] J. Howard and S. Glugger, Deep Learning for Coders with fastai and Pytorch (2020), O’Reilly Publishing
[4] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher, Wikitext (2016, Pointer Sentinel Mixture Models). This dataset is available under a Creative Commons Attribution-ShareAlike License.
