Word Embedding Revisited: A New Representation Learning and Explicit Matrix Factorization Perspective

1 Introduction

Distributed word representations, i.e., word embeddings, have been widely applied in various text mining and natural language processing tasks. Different to the traditional onehot representation that represents a word with a large vector, distributed word representation embeds every word into a low dimensional continuous space. Such representations are assumed to convey semantic and syntactic information of words.

The most widely adopted discipline to learn word embedding vectors is to maximize the corpus likelihood by neural network training [Bengio et al., 2006; Mnih and Hinton, 2007; 2009; Collobert et al., 2011; Huang et al., 2012; Le and Mikolov, 2014; Huang et al., 2013; Pennington et al., 2014; Kiros et al., 2014]. Among all these methods, CBOW and the skip-gram model in the well-known word2vec toolbox [Mikolov et al., 2013a; 2013b] have attracted tremendous attention from both academia and industry due to their effectiveness and efficiency. Remarkably, it is shown that the embedding vectors produced by these models preserve the syntactic and semantic relations between words under simple linear operations. For example, relations are captured as Madrid - Spain ≈ Paris - France, good - best ≈ great - greatest.

We focus on the skip-gram negative sampling (SGNS) model in this paper. Despite of its significant success, the underlying theoretical principles of the SGNS model are not clear enough yet. The SGNS model follows the routine of maximizing the conditional probability of the observed contexts given the current word when scanning through the corpus, however it is not clear what information the embedding vectors really convey.

The first attempt to answer this question is proposed in [Omer and Yoav, 2014], where the authors show that SGNS is an implicit matrix factorization (IMF) that factorizes an implicit word-context matrix where the value of each entry indicates the strength of association between the corresponding word-context pair. Specifically, the implicit word-context matrix that SGNS is factorizing is known as the pointwise mutual information (PMI) matrix constructed from the raw co-occurrence matrix. The authors demonstrate that in the ideal case, i.e., the embedding size of word vectors be infinity, SGNS perfectly reconstructs the PMI matrix.

In this paper, we provide a new perspective for the nature of the SGNS model. To be more specific, our main contribution is that we reformulate the objective of SGNS as a representation learning objective that has never been discovered before. In our representation learning view, the embedding vector learned for a word is a hidden representation for occurrences of the corresponding contexts under a softmax loss. Based on the representation learning view, we further show that SGNS is in fact an explicit matrix factorization (EMF), where the matrix to be factorized is the co-occurrence matrix. From this representation view, the extended task of supervised word embedding will have a clear definition. Compared with the existing IMF analysis of SGNS, our EMF formulation differs in the matrix to be factorized and the reconstruction loss. Furthermore, based on our formulation, the convergence property of the proposed algorithm will be much easier to analyze.

More importantly, once the equivalence of SGNS, representation learning and EMF is established, the new perspective will provide a solid basis for further natural extensions and generalizations of SGNS.

The remainder of the paper is organized as follows. We first review related background including the co-occurrence matrix, PMI, SGNS and general matrix factorization in Section 2, then present the perspective of representation learning and EMF for SGNS in Section 3. The optimization algorithm is designed in Section 4, followed with insight from our proposed perspective and the extended approach of supervised explicit matrix factorization. We conduct experimental investigation in Section 6 and demonstrate that the algorithm based on our formulation performs as well as the SGNS in the word2vec toolbox.

2 Background

Our explanation of word embedding focuses on skip-gram negative sampling (SGNS) from word2vec. In the following we give a brief review of the co-occurrence matrix, PMI, SGNS and general matrix factorization.

2.1 Co-occurrence Matrix and Pointwise Mutual Information

Given a training corpus D and a word w, several words are selected as the context words for w according to certain strategies, e.g., those neighboring words falling in all the fix-sized windows centered at w. We denote the number of times that a context word c appears in w's contexts as #(w, c) which is also called the co-occurrence count, and we have the following:

where VW and VC denote the word and context vocabularies in a text corpus D. The co-occurrence matrix is denoted as D, where the entry in the cth row and wth column is #(w, c). A column of D can be regarded as an explicit representation for corresponding w denoted as the explicit word vector dw in the rest of the paper. Based on D, the PMI matrix M can be constructed such that Mw,c = log( #(w,c)|D| #(w)#(c) ).

2.2 Skip-GramNegative Sampling

For a word w ∈ VW and a context word c ∈ VC, their embedding vectors are represented as column vectors w ∈ Rd and c ∈ Rd respectively, where d is the embedding's dimensionality. The embedding vectors of the words in VW and the context words in VC constitute the columns of the word and context embedding matrices W and C.

The general skip-grammodel is a simplified statistical language model that aims to predict context words given a central word w. The conditional probability is defined in a softmax form:

The context embedding matrix C and word embedding matrix W can then be learned through optimizing the following:

However, it is difficult to calculate the partition function ∑c′∈VC ewT c′ in the denominator of (1). Therefore, it is proposed in [Mikolov et al., 2013b; Mnih and Kavukcuoglu, 2013] to maximize an alternative likelihood exploiting negative sampling

where σ(x) = 1 1+e-x is the sigmoid function and PD is a probability measure taken on all words to sample false context words and k is a hyper parameter specifying the number of such words for each w. This formulation is easy to compute and is named as Skip-Gram Negative Sampling (SGNS). SGNS is applied to all the selected word-context pairs by skip-gram, leading to the following objective:

To simplify analysis, following [Omer and Yoav, 2014], PD is set to the empirical unigram distribution PD(c) = #(c) |D| .

2.3 Matrix Factorization and Representation Learning

We revisit matrix factorization and representation learning in this subsection. Given a matrix X, matrix factorization aims to factorize X into two matrices Y and Z such that X ≈ YZ. The objective of matrix factorization can be written as:

where L(·, ·) is the loss function that measures the distance between two matrices with the same size.

Different L(·, ·) leads to different matrix factorization models, including non-negative matrix factorization [Lee and Seung, 2001], probabilistic matrix factorization [Mnih and Salakhutdinov, 2007], max-margin matrix factorization [Srebro et al., 2004]. All of these models can be formulated as special cases of (5).

We can also consider matrix factorization as Representation Learning. Under certain circumstance, the objective (5) can be re-written as:

where L(·, ·) evaluates the distance between the ith column of X and YZ, and L is called the representation loss. From this point of view, one can regard that zi is the code and a hidden representation of the observed instance vector xi and Y is regarded as the representation dictionary, therefore the objective of matrix factorization can be viewed as learning a representation for each sample xi. Sparse coding [Lee et al., 2006] is a well-known representative of representation learning models.

3 SGNS as Representation Learning and Explicit Matrix Factorization

In this section, we provide a novel view for SGNS in the perspective of representation learning and explicit matrix factorization (EMF) that has never been discovered before. To be concrete, we will prove that the objective of SGNS (4) is a special case of (5) and (6), and define the specific loss function L(., .). As a consequence, we can view SGNS as a representation learning as well as a matrix factorization model. From the representation learning view, the instance vector to be represented here is the explicit word vector dw and the representation dictionary is the context embedding matrix C. From the matrix factorization view, the matrix to be factorized is the words' co-occurrence matrix D.

3.1 Representation Loss for Explicit Word Vector

Before proceeding to our main results, we define the explicit word vector, and the specific representation loss L(·, ·) in (6) for our problem.

Definition 1. For a word w, its explicit word vector dw ∈ R|VC| is defined as: dw,c = #(w, c), where dw,c is the cth element of vector dw and #(w, c) is the co-occurrence count between word w and c in the corpus.

We also define the candidate set Sw of possible explicit word vectors for word w: Sw = Sw,1 × Sw,2 × ...Sw,c × ...Sw,|Vc| is the Cartesian product of |VC| subsets and each subset is defined as Sw,c = {0, 1, · · · ,Qw,c}, where Qw,c is a pre-defined upper bound for the possible co-occurrence count between word w and c. That is, we assume that all the possible explicit word vectors for word w are elements of the set Sw. The detailed value for Qw,c is specified in the later section and we guarantee that #(w, c) ≤ Qw,c, thus dw ∈ Sw.

Given the above definition, we introduce the concept of Representation Loss for Explicit Word Vector.

Definition 2. Representation Loss for Explicit Word Vector is defined as the negative log probability of observing the explicit word vector dw given w and C. To be more concrete,

where LS serves as the representation loss and Cc denotes the cth column of C. It should be mentioned that the summation ∑d′w,c∈Sw,c is a summation in Hamming space that will be defined below in the proof of Theorem 1.

The negative representation loss LS is inspired from replicated softmax [Hinton and Salakhutdinov, 2009] in which a generative model for representing documents is proposed, and P(dw|CT w) here is the conditional distribution of the generative model in replicated softmax.

3.2 Equivalence of SGNS and EMF

Theorem 1. For a word w, when Qw,c is set to k #(w)#(c) |D| + #(w, c), LS(dw,CT w) is equivalent to -l(w,C), where l(w,C) is the loss term for w in (4).

We defer the detailed proof for Theorem 1 to the end of this subsection. To summarize, Theorem 1 guarantees that for word w, LS(dw,CT w) is equivalent to -l(w,C) in (4), if the pseudo context length Qw,c is set to an appropriate value. Therefore, the objective of SGNS (4) is equivalent to minw,C ∑w∈VW LS(dw, CT w).

According to general representation learning in Section 2.3, the equivalence shown in Theorem 1 implies that the embedding vector w learned in SGNS is a hidden representation vector for each explicit vector dw under the representation loss (7).

As the main corollary derived from Theorem 1, we point out that the objective of SGNS (4) is equivalent to explicit matrix factorization (EMF) of matrix D:

where MF denotes the loss of matrix factorization and corresponds to L(X,YZ) in (5). Each column of CT W is supposed to approximate the corresponding dw vector under the representation loss (7). In addition, from the last line of (8), we can note that MF intends to minimize an inner product loss with a soft maximum regularization, and the objective is convex over CT W, though it is not jointly convex over C and W.

Proof of Theorem 1. First the partition function of P(dw,c|w,Cc) in (7) can be factorized as follow:

so we reformulate P (dw,c|w, Cc) as:

Given the factorization of P (dw,c|w, C) above, with Cc = c and dw,c = #(w, c), LS(dw, CT w) can be reformulated as:

By substituting Qw,c = k #(w)#(c) |D| +#(w, c) into (10) above and following the recipe in [Omer and Yoav, 2014], we get

Therefore, given appropriate Qw,c, the negative representation loss LS (7) is equivalent to -l(w,C) in (4).

4 Optimization Algorithm

Given the equivalence of SGNS, representation learning and EMF established above, we propose a new optimization algorithm to train word embeddings based on the matrix factorization formulation. Specifically, we leverage the alternating minimization scheme for optimization, which is an effective and widely adopted method in the matrix factorization literature. Before we formally describe the algorithm, we derive the gradients of (8) that will be used in the algorithm:

where Ed′w |w,C denotes the conditional expectation taken on d′w with respect to the distribution P(dw|CT w) in (7) and ED′ |CT W denotes the concatenation of all Ed′w |CT w from w ∈ VW. The expectation Ed′w |CT w [d′w ] can be computed in a closed form. According to (7) and (9), we have:

where Ed′w,c |CTc w [d′w,c ] is the cth entry in the vector Ed′w |CT w [d′w ] and is the expectation of a binomial distribution given the probability σ(Ccw) and number of trials Qw,c.

The details of the Alternating Minimization for Explicit Matrix Factorization (AMEMF) algorithm are summarized in Algorithm 1. In the main loop, there are two minimization steps (starting from line 4 and line 9 respectively), which minimize the objective with regard to W and C alternatively. For each minimization step, we employ the gradient descent scheme based on the gradients in (12).

Algorithm 1: Alternating minimization for explicit matrix factorization

Input: Co-occurrence matrix D, step-size of gradient descent η, maximum number of iterations K

The main difference between AMEMF and the algorithm for SGNS in word2vec is that AMEMF is a batch alternating minimization algorithm, while word2vec scans through the corpus and updates w and c in a stochastic mode. Subtle differences in implementation include whether a negative sampling procedure is used (SGNS) or not (AMEMF); and the way to tune the learning rates: whether it is a linearly dropping learning rate to guarantee the convergence of stochastic gradient descent (SGD) in SGNS or a constant learning rate in AMEMF.

Despite of the differences, in the experiments in Section 6, we will demonstrate that the optimization algorithm AMEMF derived directly from the matrix factorization formulation (8) performs comparably with the SGD algorithm adopted in word2vec, which further verifies the equivalence between SGNS and EMF.

4.1 Convergence Analysis

Here we analyze the convergence property of the proposed algorithm AMEMF. Firstly, note that the objective in each minimization subprocedure (line 4 and 9 in Algorithm 1) is convex, which guarantees the optimal solution of each subprocedure can be reached with sublinear convergence rate [Nesterov, 2004] when the step-size is chosen properly. The whole minimization procedure can be summarized as an iteration over the following two steps: Wi = argminWMF(D,CTi-1W) and Ci = argminCMF(D,CT Wi), which implies MF(D,CTi-1Wi-1) ≥ MF(D,CTi-1Wi) ≥ MF(D,CTi Wi). As a consequence, it is guaranteed that the objective (8) descents monotonically and the algorithm will converge due to the lower bounded objective.

5 Insight from the Equivalence

We have presented the main results of this paper by establishing the equivalence between SGNS and representation learning as well as matrix factorization. Meanwhile, the contribution of this work is beyond the equivalence itself in the sense that we can now derive natural extensions of word embedding by leveraging the equivalence. To be more concrete, we propose an extended task Supervised Word Embedding, which is of great practical importance and can be conducted much more naturally within the framework of representation learning and matrix factorization than the skip-gram model. This enhances the significance of the equivalence which not only helps understanding the skip-gram model, but also benefits real world tasks.

5.1 SupervisedWord Embedding

Before describing the Supervised Word Embedding task, we first review the analogical reasoning task proposed in [Le and Mikolov, 2014], which serves as the basis for our task. The analogical reasoning task aims to find the most proper answer d in a query a - b ≈ c - d where d is not known, such as Madrid - Spain ≈ Paris - France and good - best ≈ great - greatest. Given the corresponding embedding vectors wa, wb, wc, we find the most proper word d through finding the most similar vector for wb -wa +wc under cosine distance. The percentage of these answers that we predict accurately in queries is taken as the evaluation metric for this task.

In practice, the queries in the analogical reasoning task can not only act as the benchmark for evaluating word embedding, but also provide additional side information to guide the embedding process and boost the performance. We refer this task of leveraging the information from the analogical reasoning queries to improve word embedding as Supervised Word Embedding. Although our proposed EMF is an unsupervised model as SGNS, the formulation (8) can be generalized to a supervised model naturally to incorporate supervised information based on the representation learning view, given the training corpus and analogical reasoning queries. Specifically, a query in the analogical reasoning task has four components that are denoted as "a", "b", "c" and the accurate answer "d", and the corresponding co-occurrence vectors and embedding vectors are denoted as da, db, dc, dd and wa, wb, wc, wd respectively. The objective of supervised explicit matrix factorization (SEMF) can then be written as:

where matrices Dd,Wa,Wb,Wc consist of dd, wa, wb, wc as columns respectively, and provide supervised information. λ is a hyper-parameter controlling the degree of supervision. In this model, wd is a representation for dd due to the first term and wb -wa +wc also corresponds to a representation for dd due to the second term , and thus wd is supposed to approximate wb - wa + wc. In this way, side information can be leveraged to guide the embedding process. we can observe that the proposed supervised model is solely based on our MF matrix factorization loss and is a natural extension of EMF.

The supervised model (14) can be solved similarly with the alternating minimization framework as in Algorithm 1.

6 Experiments

In this section, we conduct several experiments to verify the effectiveness of the AMEMF algorithm in the EMF framework. The experiments consist of two parts: comparison of word embedding methods and evaluation of supervised word embedding. They are both evaluated by the analogical reasoning task. To compare the methods, we will evaluate the performance of SGNS, SPPMI, IMF and EMF, and verify the equivalence of SGNS and EMF. Here SPPMI indicates the Shifted Positive Pointwise Mutual Information (PMI) matrix MSPPMI , which is a variant of PMI as MSPPMI w,c = max(Mw,c - log k, 0) and has better performance than the original PMI in the analogical reasoning task according to [Omer and Yoav, 2014]. IMF here represents the word embedding approach based on Singular Value Decomposition (SVD) of the SPPMI matrix. In the supervised word embedding task, we will compare EMF and SEMF, and justify the superiority of SEMF by leveraging side information.

Datasets

We use a publicly accessible dataset Enwik9 1 as our training corpus. Enwik9 contains about 124 million tokens. We adopt the original dataset of analogical reasoning queries used in [Mikolov et al., 2013a] for the analogical reasoning task, which contains 19544 queries called Google query dataset. Vocabulary size is controlled by a hyper-parameter mini-count that filters low frequency (less than mini-count) words out and different vocabulary sizes result in different sizes of query dataset correspondingly. The negative sampling parameter is denoted as k as in (4).

1 http://mattmahoney.net/dc/textdata.html

Experimental Setup

The dimensionality of all embedding vectors is set to 200. For comparison of word embedding methods, we compare the embedding vectors produced by SGNS, SPPMI, IMF with SVD and EMF with the AMEMF algorithm under different k and mini-count values. The step-size of AMEMF is set to 6e - 7. To keep the settings of AMEMF and SGNS as consistent as possible, they use the same co-occurrence matrix produced by the skip-gram strategy with window size 5. The step-size of SGNS is set according to the default in the word2vec toolbox, PD is set to the unigram distribution, and the number of iterations K in AMEMF is set to 200.

For the supervised word embedding task, we partition the Google query dataset into a training query set and a test query set, and then train the supervised model (SEMF) on the training query set and test its embedding vectors in the analogical reasoning task. We compare SEMF and the unsupervised model (EMF) with different K values and training ratios. The mini-count value is fixed to 1000, the step-size is set to 5e-7, and the number of iterations K in AMEMF is set to 100.

6.1 Experimental Results and Analysis

We evaluate SGNS, SPPMI, IMF and EMF in the analogical reasoning task. The comparison of these four models in terms of the analogical reasoning accuracy is shown in Table 1 with different values of the negative sampling parameter k and mini-count. For each cell corresponding to a combination of k and mini-count values, we highlight the best performance. The SPPMI model achieves its best performance when k is set to 1 and the its accuracy decreases rapidly when k increases, therefore, we evaluate it independently. With k = 1, one can observe that the accuracy of SPPMI is not comparable to the best performance in Table 1.

In the rest 12 cells of Table 1, EMF performs best in 8 cells, while SGNS performs best in 5 cells. Though it appears that EMF is subtly better than SGNS, the accuracy of EMF and SGNS are very close actually. Generally speaking, EMF and SGNS perform quite similarly for all the different k and minicount parameter pairs and are significantly superior to the other models. This experimental result empirically serves as a strong evidence of the equivalence between EMF and SGNS. Meanwhile, the IMF model does not perform well in the analogical reasoning task which matches the experimental result from [Omer and Yoav, 2014], and this also demonstrates that EMF could be a more reasonable connection between SGNS and matrix factorization than IMF.

Next the unsupervised model EMF and supervised model SEMF are evaluated with analogical reasoning test accuracy in Table 2. There is no natural way to incorporate supervised information for the IMF model, which is another advantage of EMF over IMF. Table 2 shows that SEMF significantly outperforms EMF, verifying that incorporating supervised information improves the quality of word embeddings. Given the similar performance of EMF and SGNS in the analogical reasoning task, the superiority of SEMF over the original word2vec is obvious. As to the effects of different hyper parameters, one can observe that the performance of the unsupervised EMF remains stable with different training ratios. In the meantime, it can be observed that the performance of SEMF grows with the training ratio and λ impacts the degree of supervision significantly. This experiment not only verifies the generalization ability of SEMF, but also inspires that we can equip embedding vectors with intended properties through supervision.

7 Conclusion

We revisit the skip-gram negative sampling (SGNS) model in the popular toolbox word2vec, and prove that intrinsically SGNS is a representation learning method, as well as an explicit matrix factorization (EMF) of the co-occurrence matrix that is directly obtained from corpus. Different to implicit matrix factorization (IMF) [Omer and Yoav, 2014], our objective is explicitly equivalent to SGNS, based on which the task of supervised word embedding can be conducted naturally. Experimental results justify the equivalence between SGNS and EMF, as well as the validity of supervised word embedding.

Acknowledgments

We would like to thank Prof. Qing Ling for his advice on algorithm design. This work is supported by the National Natural Science Foundation of China (No. 61375060), and the National Science Foundation for Distinguished Young Scholars of China (No. 61325010).