CBoW PyTorch

Training CBoW Model

This notebooks is a part of AI for Beginners Curriculum

In this example, we will look at training CBoW language model to get our own Word2Vec embedding space. We will use AG News dataset as the source of text.

First let's load our dataset and define tokenizer and vocabulary. We will set vocab_size to 5000 to limit computations a bit.

CBoW Model

CBoW learns to predict a word based on the $2N$ neighboring words. For example, when $N=1$, we will get the following pairs from the sentence I like to train networks: (like,I), (I, like), (to, like), (like,to), (train,to), (to, train), (networks, train), (train,networks). Here, first word is the neighboring word used as an input, and second word is the one we are predicting.

To build a network to predict next word, we will need to supply neighboring word as input, and get word number as output. The architecture of CBoW network is the following:

• Input word is passed through the embedding layer. This very embedding layer would be our Word2Vec embedding, thus we will define it separately as embedder variable. We will use embedding size = 30 in this example, even though you might want to experiment with higher dimensions (real word2vec has 300)
• Embedding vector would then be passed to a linear layer that will predict output word. Thus it has the vocab_size neurons.

For the output, if we use CrossEntropyLoss as loss function, we would also have to provide just word numbers as expected results, without one-hot encoding.

Preparing Training Data

Now let's program the main function that will compute CBoW word pairs from text. This function will allow us to specify window size, and will return a set of pairs - input and output word. Note that this function can be used on words, as well as on vectors/tensors - which will allow us to encode the text, before passing it to to_cbow function.

Let's prepare the training dataset. We will go through all news, call to_cbow to get the list of word pairs, and add those pairs to X and Y. For the sake of time, we will only consider first 10k news items - you can easily remove the limitation in case you have more time to wait, and want to get better embeddings :)

We will also convert that data to one dataset, and create dataloader:

We will also convert that data to one dataset, and create dataloader:

Now let's do the actual training. We will use SGD optimizer with pretty high learning rate. You can also try playing around with other optimizers, such as Adam. We will train for 10 epochs to begin with - and you can re-run this cell if you want even lower loss.

Trying out Word2Vec

To use Word2Vec, let's extract vectors corresponding to all words in our vocabulary:

Let's see, for example, how the word Paris is encoded into a vector:

It is interesting to use Word2Vec to look for synonyms. The following function will return n closest words to a given input. To find them, we compute the norm of $|w_i - v|$, where $v$ is the vector corresponding to our input word, and $w_i$ is the encoding of $i$-th word in the vocabulary. We then sort the array and return corresponding indices using argsort, and take first n elements of the list, which encode positions of closest words in the vocabulary.

Takeaway

Using clever techniques such as CBoW, we can train Word2Vec model. You may also try to train skip-gram model that is trained to predict the neighboring word given the central one, and see how well it performs.