Implementing word2vec in PyTorch (skip-gram model) | by Mateusz Bednarski | Towards Data Science

You probably have heard about word2vec embedding. But do you really understand how it works? I though I do. But I have not, until implemented it.
This is why I’m creating this guide.
2021 Update: For a more detailed article see: https://neptune.ai/blog/word-embeddings-guide
I assume you know more-less what word2vec is.
Corpus
In order to be able to track every single step i’m using following nano corpus:
Very first step is word2vec to create the vocabulary. It has to be built at the beginning, as extending it is not supported.
Vocabulary is basically a list of unique words with assigned indices.
Corpus is very simple and short. In real implementation we would have to perform case normalization, removing some punctuation etc, but for simplicity let’s use this nice and clean data. Anyway, we have to tokenize it:
This will give us a list of tokens:
[['he', 'is', 'a', 'king'], ['she', 'is', 'a', 'queen'], ['he', 'is', 'a', 'man'], ['she', 'is', 'a', 'woman'], ['warsaw', 'is', 'poland', 'capital'], ['berlin', 'is', 'germany', 'capital'], ['paris', 'is', 'france', 'capital']]
We iterate over tokens in corpus, and generate list of unique words(tokens). Next we create two dictionaries for mapping between word and index.
Which gives us:
 0: 'he', 1: 'is', 2: 'a', 3: 'king', 4: 'she', 5: 'queen', 6: 'man', 7: 'woman', 8: 'warsaw', 9: 'poland', 10: 'capital', 11: 'berlin', 12: 'germany', 13: 'paris', 14: 'france'
We can now generate pairs center word, context word. Let’s assume context window to be symmetric and equal to 2.
It gives us list of pairs center, context indices:
array([[ 0,  1],       [ 0,  2],       ...
Which can be easily translated to words:
he ishe ais heis ais kinga hea isa king
Which makes perfect sense.
Now, we are going through details from very first equation to working implementation.
For skip-gram we are interested in predicting context, given center word and some parametrization. This is our probability distribution for single pair.
Now, we want to maximize it trough all word/context pairs.
Wait, why?
As we are interested in predicting context given center word, we want to maximize P(context|center) for each context, center pair. As probability sums up to 1 — we are implicitly making P(context|center) close to 0 for all non-existing context, center pairs. By multiplying those probabilities we make this function close to 1 if our model is good and close to 0 if it is bad. Of course we are pursuing a good one — so there is max operator at the beginning.
This expression is not very suitable to compute. That’s why we are going to perform some very common transformations.
Step 1 — Replace probability with negative log likelihood.
Recall that neural nets are about minimizing loss function. We could simply multiply P by minus one , but applying log gives us better computational properties. This does not change position of function extrema (as log is a strictly monotonic function). So expression is changed to:
Step 2 — Replace products with sums
Next step is to replace products with sums. We can do it because:
Step 3 — Transform to a proper loss function
And after dividing by number of pars (T) we have our final loss term:
Great, but how we do define P(context|center)? For now, let’s assume that reach word has actually two vectors. One if is present as center word (v), and second one if context(u). Given that definition for P looks following:
That’s scary!
Let me break it down to smaller parts. See following structure:
This is just a softmax function. Now closer look at nominator
Both u and v are vectors. This expression is just scalar product of given center, context pair. Bigger as they are more similar to each other.
Now, denominator:
We’re iterating over all words in vocabulary.
And computing “similarity” for given center word and each word in vocabulary treated as context word.
To sum up:
For each existing center, context pair in corpus we’re computing their “similarity score”. And divide it by sum of each theoretically possible context — to know whether score is relatively high or low. As softmax is guaranteed to take a value between 0 and 1 it defines a valid probability distribution.
Neural net implementing this concept consists of three layers: input, hidden and output.
Input layer is just the center word encoded in one-hot manner. It dimensions are [1, vocabulary_size]
Hidden layer makes our v vectors. Therefore it has to have embedding_dims neurons. To compute it value we have to define W1 weight matrix. Of course its has to be [embedding_dims, vocabulary_size]. There is no activation function — just plain matrix multiplication.
What’s important — each column of W1 stores v vector for single word. Why? Because x is one-hot and if you multiply one-hot vector by matrix, result is same as selecting select single column from it. Try on your own using a piece of paper ;)
Last layer must have vocabulary_size neurons — because it generates probabilities for each word. Therefore, W2 is [vocabulary_size, embedding_dims] in terms of shape.
On top on that we have to use softmax layer. PyTorch provides optimized version of this, combined with log — because regular softmax is not really numerically stable:
log_softmax = F.log_softmax(a2, dim=0)
This is equivalent to compute softmax and after that applying log.
Now we can compute loss. As usual PyTorch provides everything we need:
loss = F.nll_loss(log_softmax.view(1,-1), y_true)
The nll_loss computes negative-log-likelihood on logsoftmax. y_true is context word — we want to make this as high as possible — because pair x, y_true is from training data — so the are indeed center, context.
As we fished forward pass, now it’s time to perform backward pass. Simply:
loss.backward()
For optimization SDG is used. It is so simple, that it was faster to write it by hand instead of creating optimizer object:
W1.data -= 0.01 * W1.grad.dataW2.data -= 0.01 * W2.grad.data
Last step is to zero gradients to make next pass clear:
W1.grad.data.zero_()W2.grad.data.zero_()
Time to compile it into training loop. It can look like:
One potentially tricky thing is y_true definition. We do not create one-hot explicitly — nll_loss does it by itself.
Loss at epo 0: 4.241989389487675Loss at epo 10: 3.8398486052240646Loss at epo 20: 3.5548086541039603Loss at epo 30: 3.343840673991612Loss at epo 40: 3.183084646293095Loss at epo 50: 3.05673006943294Loss at epo 60: 2.953996729850769Loss at epo 70: 2.867735825266157Loss at epo 80: 2.79331214427948Loss at epo 90: 2.727727291413716Loss at epo 100: 2.6690095041479385
Ok, we have trained the network. One, last thing is to extract vectors for words. It is possible in three ways:
Use vector v from W1
Use vector u from W2
Use average v and u
Try to think on your own when to use which one ;)
I’m working on online interactive demonstration on this. It should be available soon. Stay tuned ;)
You can download code from here.