Incredible Word2Vec: 5. TensorFlow Version of Word2Vec

Differences

For the basic mathematical principles of Word2Vec, please refer to the article "Incredible Word2Vec: 1. Mathematical Principles". The main models in this article are still CBOW or Skip-Gram, but the loss design is different. This article still uses a full softmax structure rather than Hierarchical Softmax or Negative Sampling schemes. However, when training the softmax, a cross-entropy loss based on random negative sampling is used. This loss is different from the existing nce_loss and sampled_softmax_loss; for now, let’s name it random softmax loss.

Additionally, in a softmax structure, the form is generally \text{softmax}(Wx+b). Considering that the shape of the W matrix is actually the same as the shape of the word vector matrix, this article considers a model where the softmax layer and the word vector layer share weights (in this case, b is set directly to 0). This model is equivalent to the original Word2Vec negative sampling scheme and is similar to the word co-occurrence matrix factorization of GloVe word vectors. However, because cross-entropy loss is used, it theoretically converges faster, and the training results still retain the predictive probability meaning of softmax. (In contrast, after training existing Word2Vec negative sampling models, the final output values of the model are meaningless; only the word vectors are meaningful.) At the same time, because parameters are shared, the word vectors are updated more thoroughly. Readers are encouraged to test this scheme.

Therefore, this article actually implements four model combinations: CBOW/Skip-Gram, with or without shared softmax layer parameters. Readers can choose which to use.

Origin of the Loss

As mentioned earlier, one of the main purposes of this article is to test the effectiveness of the new loss. Below is a brief introduction to the origin and form of this loss. This starts with why softmax is difficult to train.

Assume the number of labels (which in this article is 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} Here Z = e^{z_1} + e^{z_2} + \dots + e^{z_n}. If the correct category label is t, using cross-entropy as the loss gives: L = -\log \frac{e^{z_t}}{Z} The gradient is: \nabla L = -\nabla z_t + \nabla (\log Z) = -\nabla z_t + \frac{\nabla Z}{Z} Because of the existence of Z, every time gradient descent is performed, the complete Z must be calculated to compute \nabla Z. This means the computational complexity for a single sample iteration is \mathcal{O}(n). For cases where n is large, this is unacceptable, so approximation schemes are sought. (Hierarchical softmax is one such scheme, but it is complex to implement, its results are usually slightly worse than ordinary softmax, and while it speeds up training, it is actually slower if you need to find the label with the maximum probability during prediction.)

Let’s further calculate \nabla L: \begin{aligned} \nabla L &= -\nabla z_t + \frac{\sum_i e^{z_i} \nabla z_i}{Z} \\ &= -\nabla z_t + \sum_i \frac{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, and the other is the mean of the gradients of all labels. These two terms have opposite signs and can be understood as a "tug-of-war." The computational load is mainly concentrated in the second term because it requires traversing all labels to calculate the specific mean. However, since the mean itself has probabilistic significance, can we simply randomly select a few labels to calculate this gradient mean instead of calculating all gradients? If so, the computational load per update step would be fixed and would not increase rapidly as the number of labels grows.

However, doing this requires randomly selecting labels according to their probabilities, which is not easy to code. A more ingenious method is to avoid calculating the gradient directly and instead modify the loss function. This leads to the loss used in this article: For each "sample-label" pair, randomly select nb_negative labels, combine them with the original label to form nb_negative + 1 labels, and calculate the softmax and cross-entropy directly within these labels. After selecting such a loss and calculating the gradient, you will find that it naturally becomes the gradient mean selected according to probability.

Code Implementation

I feel the code is quite concise, with a single file containing everything. The training output mimics Word2Vec in gensim. The model code is located 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']

# Build and train the model
wv = Word2Vec(texts(), model='cbow', nb_negative=16, shared_softmax=True, epochs=2) 
# Save the model to the 'myvec' folder in the current directory
wv.save_model('myvec') 

# After training, you can call it like this:
wv = Word2Vec() # Create an empty model
wv.load_model('myvec') # Load the model from the 'myvec' folder

A few points to clarify:

1. The training input consists of segmented sentences, which can be a list or an iterator (class + __iter__). Note that it cannot be a generator (function + yield), which is consistent with the requirements of the gensim version of Word2Vec. This is because a generator can only be traversed once, while training Word2Vec requires multiple passes over the data.

2. The model does not support incremental training; that is, once the model is trained, it cannot be updated with additional documents. (It’s not impossible, but it’s unnecessary and doesn’t add much value.)

3. Training the model requires TensorFlow, and GPU acceleration is recommended. After training is complete, reloading and using the model does not require TensorFlow, only NumPy.

4. Regarding the number of iterations, 1 to 2 epochs are usually sufficient, and 10 to 30 negative samples are enough. Other parameters, such as batch_size, can be adjusted through experimentation.

Simple Comparative Experiments

In TensorFlow, the two existing losses for approximating softmax training are nce_loss and sampled_softmax_loss. Here is a simple comparison. We trained the same model on a travel domain corpus (over 20,000 articles) and compared the results. The model used was CBOW, the softmax layer did not share weights with the word vector layer, and all other parameters used the same default values.

Future Work

One question is: why create a new loss if it’s comparable to sampled_softmax? The reason is simple: when I look at the paper and formulas for sampled_softmax, I always feel it’s not quite "beautiful" and the theory isn’t elegant enough. Of course, based on the results, I’m just being too much of a perfectionist. This article can be considered a product of that perfectionism, as well as a way to practice TensorFlow.

Additionally, the article "Recording a Semi-supervised Sentiment Analysis" shows that language models have significant potential in pre-training models and implementing semi-supervised tasks. Even word vectors are just the first layer of parameters pre-trained by a language model. Therefore, I want to take some time to delve deeper into similar content. This article is one of the preparations for that.

When reposting, please include the original address: https://kexue.fm/archives/4402

For more detailed reposting matters, please refer to: Scientific Space FAQ