【The Amazing Word2Vec】5. Word2Vec in Tensorflow

【不可思议的Word2Vec】5. Tensorflow版的Word2Vec 苏剑林 · 2017-05-27 · Machine Learning deep learning word vector word2vec

This post packages a fairly complete implementation of Word2Vec, with the model part implemented in Tensorflow. The purpose of this post is not merely to reinvent the Word2Vec wheel yet again, but to use this example as a way to get familiar with Tensorflow's coding style, and to test the effectiveness of a new softmax loss designed by the author, in preparation for later work on language models.

What's different here

For the basic mathematical principles of Word2Vec, please head over to 《【The Amazing Word2Vec】1. Mathematical Principles》. The main model here is still CBOW or Skip-Gram, but the loss is designed differently. This post still uses the full softmax structure, rather than the hierarchical (Huffman) softmax or negative-sampling schemes, but when training the softmax, it uses a cross-entropy loss based on random negative sampling. This loss is different from both the existing nce_loss and sampled_softmax_loss; here we tentatively call it "random softmax loss."

In addition, in the softmax structure, the usual form is $\text{softmax}(Wx+b)$. Considering that the shape of the $W$ matrix is in fact identical to that of the word-embedding matrix, this post also considers a model in which the softmax layer shares weights with the embedding layer (in which case $b$ is simply set to zero). This model is equivalent to the original Word2Vec negative-sampling scheme, and is also similar to GloVe's factorization of the word co-occurrence matrix. But since it uses cross-entropy loss, it should in theory converge faster, and the trained output still retains the probabilistic meaning of a softmax (by contrast, once the existing Word2Vec negative-sampling model has been trained, the model's output values themselves are meaningless — only the word vectors are meaningful). Moreover, because the parameters are shared, the word vectors get updated more thoroughly; readers are encouraged to experiment further with this scheme. more

So, in effect, this post implements 4 combinations of models: CBOW/Skip-Gram, with/without sharing the softmax layer's parameters — readers can pick whichever they like.

Where does the loss come from?

As mentioned above, one of the main purposes of this post is to test the effectiveness of a new loss. Let's briefly go over where it comes from and what form it takes. This starts with the question of why softmax is hard to train in the first place.

Suppose the number of labels (in this post, the vocabulary size) is $n$. Then

$$\begin{aligned}(p_1,p_2,\dots,p_n) =& \text{softmax}(z_1,z_2,\dots,z_n)\\ =& \left(\frac{e^{z_1}}{Z}, \frac{e^{z_2}}{Z}, \dots, \frac{e^{z_n}}{Z}\right)\end{aligned}$$

where $Z = e^{z_1} + e^{z_2} + \dots + e^{z_n}$. If the correct class label is $t$, and we use cross-entropy as the loss, then

$$L=-\log \frac{e^{z_t}}{Z}$$

with gradient

$$\nabla L=-\nabla z_t + \nabla (\log Z)=-\nabla z_t + \frac{\nabla Z}{Z}$$

Because of the presence of $Z$, every gradient-descent step requires computing the full $Z$ in order to compute $\nabla Z$ — that is, the computational cost of a single sample's update is already $\mathcal{O}(n)$. For large $n$, this is unacceptable, so we need to look for approximation schemes (Huffman softmax is one such approach, but it's rather complicated to implement, and its results are usually slightly worse than plain softmax; moreover, Huffman softmax only speeds up training — if at inference time you need to find the label with the highest probability, it's actually slower).

Let's take the computation of $\nabla L$ a bit further:

$$\begin{aligned}\nabla L=&-\nabla z_t + \frac{\sum_i e^{z_i}\nabla z_i}{Z}\\ =&-\nabla z_t + \frac{\sum_i e^{z_i}}{Z}\nabla z_i\\ =&-\nabla z_t + \sum_i p_i \nabla z_i\\ =&-\nabla z_t + \text{E}(\nabla z_i) \end{aligned}$$

In other words, the final gradient consists of two terms: one is the gradient of the correct label, the other is the average gradient over all labels, and these two terms have opposite signs — we can think of them as being in a "tug of war." Most of the computational cost is concentrated in the second term, since we need to iterate over every label to compute this exact average. However, this average itself has a probabilistic interpretation — so could we just randomly sample a handful of labels to estimate this gradient average, rather than computing the gradient over all of them? If we could, then the computational cost per update step would be fixed, and would not grow rapidly with the number of labels.

But doing this directly would require sampling labels according to their probabilities, which is also not easy to implement. However, there's a cleverer approach: instead of computing the gradient ourselves directly, we can tamper with the loss instead. This is exactly what leads to the loss used in this post: for each "sample-label" pair, randomly select nb_negative labels, then combine them with the original label to form nb_negative + 1 labels, and directly compute the softmax and cross-entropy over these nb_negative + 1 labels. Once we adopt this loss and then compute its gradient, we find that it naturally turns out to be the probability-weighted average gradient we wanted.

Code implementation

I feel the code turned out to be fairly compact — the whole thing fits in a single file. The training output mimics that of gensim's Word2Vec. The model code is on GitHub:

https://github.com/bojone/tf_word2vec/blob/master/Word2Vec.py

Usage reference:

from Word2Vec import *
import pymongo
db = pymongo.MongoClient().travel.articles
class texts:
    def __iter__(self):
        for t in db.find().limit(30000):
            yield t['words']

wv = Word2Vec(texts(), model='cbow', nb_negative=16, shared_softmax=True, epochs=2) #建立并训练模型
wv.save_model('myvec') #保存到当前目录下的myvec文件夹

#训练完成后可以这样调用
wv = Word2Vec() #建立空模型
wv.load_model('myvec') #从当前目录下的myvec文件夹加载模型

A few things worth noting:

1. The training input is a set of pre-tokenized sentences, which can be a list or an iterator (a class with __iter__); note that it must NOT be a generator (a function using yield), consistent with the requirement in gensim's version of word2vec. This is because a generator can only be traversed once, while training word2vec requires traversing the data multiple times.
2. The model does not support incremental/online training — that is, once training is complete, you cannot update the existing model with additional documents (not that it's technically impossible, it's just unnecessary and of limited value).
3. Training the model requires Tensorflow, and using a GPU for acceleration is recommended. Once training is complete, reloading and using the model requires only numpy, not Tensorflow.
4. As for the number of epochs, 1–2 is generally enough, and 10–30 negative samples suffice. Other parameters like batch_size can be tuned through your own experiments.

A simple comparative experiment

In Tensorflow, the two existing approximate softmax training losses are nce_loss and sampled_softmax_loss; here we make a simple comparison. We train the same model on a travel-domain corpus (over twenty thousand articles) and compare the results. The model is CBOW, the softmax layer does not share parameters with the embedding layer, and all other parameters use the same defaults.

random_softmax_loss

Time taken: 8 minutes 19 seconds (2 epochs, batch_size 8000)

Similarity test results

import pandas as pd
pd.Series(wv.most_similar(u'fruit'))
0 (food, 0.767908)
1 (dried fish, 0.762363)
2 (coconut, 0.750326)
3 (beverage, 0.722811)
4 (foodstuff, 0.719381)
5 (beef jerky, 0.715441)
6 (pineapple, 0.715354)
7 (sausage, 0.714509)
8 (jackfruit, 0.712546)
9 (raisins, 0.709274)
dtype: object
pd.Series(wv.most_similar(u'nature'))
0 (humanities, 0.645445)
1 (harmony, 0.634387)
2 (inclusiveness, 0.61829)
3 (great nature, 0.601749)
4 (natural environment, 0.588165)
5 (blend, 0.579027)
6 (broad-minded, 0.574943)
7 (interpretation, 0.550352)
8 (wildness, 0.548001)
9 (rustic charm, 0.545887)
dtype: object
pd.Series(wv.most_similar(u'Guangzhou'))
0 (Shanghai, 0.749281)
1 (Wuhan, 0.730211)
2 (Shenzhen, 0.703333)
3 (Changsha, 0.683243)
4 (Fuzhou, 0.68216)
5 (Hefei, 0.673027)
6 (Beijing, 0.669859)
7 (Chongqing, 0.653501)
8 (Haikou, 0.647563)
9 (Tianjin, 0.642161)
dtype: object
pd.Series(wv.most_similar(u'scenery'))
0 (scenery(alt), 0.825557)
1 (beautiful view, 0.763399)
2 (landscape, 0.734687)
3 (scenic view, 0.727672)
4 (vista, 0.57638)
5 (picturesque lakes and mountains, 0.573512)
6 (mountain view, 0.555502)
7 (breathtaking, 0.552739)
8 (Mingshi (place name), 0.535922)
9 (along the way, 0.53485)
dtype: object
pd.Series(wv.most_similar(u'restaurant (酒楼)'))
0 (restaurant (酒家), 0.768179)
1 (food stall, 0.731749)
2 (hotpot restaurant, 0.729214)
3 (street stall, 0.726048)
4 (eatery, 0.722667)
5 (noodle shop, 0.715188)
6 (roadside stall, 0.709883)
7 (renowned shop, 0.708996)
8 (Songhelou (restaurant name), 0.705759)
9 (branch (store), 0.705749)
dtype: object
pd.Series(wv.most_similar(u'hotel'))
0 (Marriott, 0.722409)
1 (Hilton, 0.713292)
2 (five-star, 0.697638)
3 (five-star (alt), 0.696659)
4 (Kempinski, 0.694978)
5 (Yintai, 0.693179)
6 (grand hotel, 0.692239)
7 (guesthouse, 0.67907)
8 (Sheraton, 0.668638)
9 (holiday, 0.662169)
dtype: object

nce_loss

Time taken: 4 minutes (2 epochs, batch_size 8000); however, the similarity test results were simply unbearable. Granted, given that the training time was shorter, in fairness we increased the epochs to 4, keeping everything else the same, and ran it again. The similarity test results were still a mess, e.g.:

pd.Series(wv.most_similar(u'fruit'))
0 (mouth, 0.940704)
1 (can, 0.940106)
2 (100, 0.939276)
3 (change, 0.938824)
4 (second, 0.938155)
5 (:, 0.938088)
6 (see, 0.937939)
7 (not good, 0.937616)
8 (and, 0.937535)
9 ((, 0.937383)
dtype: object

I began to suspect I was using it incorrectly, so I adjusted things again, increasing nb_negative to 1000, and setting the epochs back to 3; this took 9 minutes 17 seconds. The final loss was an order of magnitude smaller than before, and the similarity results were somewhat more reasonable, though still not particularly good, e.g.:

pd.Series(wv.most_similar(u'fruit'))
0 (local specialty, 0.984775)
1 (seafood, 0.981409)
2 (and the like, 0.981158)
3 (food, 0.980803)
4 (., 0.980371)
5 (vegetables, 0.979822)
6 (&, 0.979713)
7 (mango, 0.979599)
8 (can, 0.979486)
9 (such as, 0.978958)
dtype: object
pd.Series(wv.most_similar(u'nature'))
0 (with/and, 0.985322)
1 (located in, 0.984874)
2 (these, 0.983769)
3 (madam, 0.983499)
4 (in/inside, 0.983473)
5 (的 [possessive particle], 0.983456)
6 (will/shall, 0.983432)
7 (former residence, 0.983328)
8 (those, 0.983089)
9 (here, 0.983046)
dtype: object

sampled_softmax_loss

With the previous experience in hand, this time we directly set nb_negative to 1000, with 3 epochs, taking 8 minutes 38 seconds. The similarity comparison results were:

pd.Series(wv.most_similar(u'fruit'))
0 (snacks, 0.69762)
1 (food, 0.651911)
2 (chocolate, 0.64101)
3 (grapes, 0.636065)
4 (biscuits, 0.62631)
5 (bread, 0.613488)
6 (cantaloupe, 0.604927)
7 (foodstuff, 0.602576)
8 (dried goods, 0.601015)
9 (pineapple, 0.598993)
dtype: object
pd.Series(wv.most_similar(u'nature'))
0 (humanities, 0.577503)
1 (great nature, 0.537344)
2 (landscape, 0.526281)
3 (pastoral, 0.526062)
4 (unique, 0.526009)
5 (harmony, 0.503326)
6 (charming, 0.498782)
7 (boundless, 0.491521)
8 (gorgeous, 0.482407)
9 (a scene of, 0.479687)
dtype: object
pd.Series(wv.most_similar(u'Guangzhou'))
0 (Shenzhen, 0.771525)
1 (Shanghai, 0.739744)
2 (Dongguan, 0.726057)
3 (Shenyang, 0.687548)
4 (Fuzhou, 0.654641)
5 (Beijing, 0.650491)
6 (EMU train, 0.644898)
7 (take the EMU train, 0.635638)
8 (Haikou, 0.631551)
9 (Changchun, 0.628518)
dtype: object
pd.Series(wv.most_similar(u'scenery'))
0 (scenery(alt), 0.8393)
1 (landscape, 0.731151)
2 (scenic view, 0.730255)
3 (beautiful view, 0.666185)
4 (snow scene, 0.554452)
5 (vista, 0.530444)
6 (picturesque lakes and mountains, 0.529671)
7 (mountain view, 0.511195)
8 (road conditions, 0.490073)
9 (picturesque, 0.483742)
dtype: object
pd.Series(wv.most_similar(u'restaurant (酒楼)'))
0 (restaurant (酒家), 0.766124)
1 (eatery, 0.687775)
2 (dining hall, 0.666957)
3 (restaurant (饭店), 0.664034)
4 (Sichuan-style, 0.659254)
5 (eatery (alt), 0.658057)
6 (food stall, 0.656883)
7 (simple homely fare, 0.650861)
8 (Gonghechun (restaurant name), 0.650256)
9 (eatery (alt2), 0.644265)
dtype: object
pd.Series(wv.most_similar(u'hotel'))
0 (guesthouse, 0.685888)
1 (grand hotel, 0.678389)
2 (four-star, 0.638032)
3 (five-star, 0.633661)
4 (Hanting, 0.619405)
5 (Home Inn, 0.614918)
6 (lobby, 0.612269)
7 (resort, 0.610618)
8 (four-star (alt), 0.609796)
9 (Tianyu, 0.598987)
dtype: object

Summary

This experiment isn't especially rigorous, but I think it's fair to say that, given the same training time and judging from the similarity task, random softmax and sampled softmax seem to perform comparably, while nce loss performs the worst. Further shrinking the number of epochs and tuning parameters showed similar results — readers are welcome to test further. Because the random softmax in this post resamples differently for every sample, fewer negative samples are needed overall, and the sampling is more thorough.

As for comparisons on other tasks, that will have to wait for future practice. After all, this isn't for publishing a paper — I'm too lazy to go further with it for now.

Follow-up work

One might ask: if it performs about the same as sampled softmax, why bother inventing a new loss? The reason is actually simple: looking at the sampled softmax paper and its formulas, I always felt it wasn't quite elegant, that the theory wasn't as clean as it could be. Of course, judging purely by results, maybe I'm just being overly obsessive. This post is, if anything, a product of that obsession, and also served as practice with Tensorflow.

Also, 《Notes on a Semi-Supervised Sentiment Analysis Experiment》 showed that language models hold great potential for pretraining, semi-supervised learning, and related tasks — indeed, word embeddings are really nothing more than the first layer of parameters obtained by pretraining a language model. So I'd like to find time to dig deeper into similar topics. This post is one of the preparations for that.