Fast Extraction of Word Embedding from Q-contexts

We first introduce the $\ell^{2}$-norm sampling, then describe the definition of Q-contexts.

$\ell^{2}$-norm sampling: The $\ell^{2}$-norm sampling technique has well exhibited its effectiveness in machine learning (Hazan et al., 2011; Song et al., 2016) and randomized linear algebra (Drineas et al., 2008). In fact, the work by Frieze, Kannan, and Vempala (Frieze et al., 2004) shows that with certain $\ell^{2}$-norm sampling assumptions, a form of singular value estimation can be achieved in time independent of the size of input matrix. Further, the work by Tang (Tang, 2019) shows that sampling from the projection of a vector onto a subspace is not outside the realm of feasibility.

Inspired by these, we leverage the $\ell^{2}$-norm sampling to construct a small typical information matrix to solve the large-scale word embedding learning problem. We now elaborate on our proposed relation between Q-contexts and words in Equation (1).

Given an information matrix $M$, we first encode the information matrix $M$ into an $\ell^{2}$-norm state which will be described in detail in Section 3.2. Now that the mutual information $M$ is prepared to be a state, each $\ell^{2}$-norm sampling yields a context $c_{i}$ with probability $p_{c_{i}}$. The collection of the corresponding row vectors $r_{c_{i}}=M_{c_{i},*}$ is what we call Q-contexts. In our proposed scheme, $p_{c}$ is designed to correctly reflect the amount of information carried by the context $c$, so that contexts with more information are more likely to be sampled.

The central idea in Equation (1) is that a few $\ell^{2}$-norm sampling of the information matrix provides sufficient information such that a good word embedding can be obtained from a linear combination of its mutual information with Q-contexts

where the context coefficients vectors (i.e., $\lambda_{c}$’s) are to be determined in Section 3.2.

We now give the full algorithm (Algorithm 2) for computing word embedding based on Q-contexts matrix $R$. WEQ method consists of three steps: information matrix calculation, Q-contexts construction, and calculating word embedding from Q-contexts.

Step 1: Information matrix calculation. As we have discussed in Section 2, there are different kinds of information matrices. Here, we choose PPMI and SPPMI matrices proposed in (Levy and Goldberg, 2014) which shows that exact factorizing (PPMI / SPPMI) matrix with SVD is at least as good as SGNS’s solutions. Given a co-occurrence matrix, we construct the PPMI information matrix as follows:

and a shifted version of PPMI (SPPMI with shift parameter ${\kappa}$)

Step 2: Q-contexts construction. As in Definition 3.1, we will encode the information matrix $M$ into an $\ell^{2}$-norm state, which admits our fast construction of Q-contexts matrix $R$.

In our method, we propose a data structure that achieves fast $\ell^{2}$-norm sampling in practice. For the information matrix $M\in\mathbb{R}^{n\times n}$, we store the cumulative summation of its row-squared $(\|r_{1}\|^{2},\|r_{1}\|^{2}+\|r_{2}\|^{2},\cdots,\|M\|_{F}^{2})$. To perform the $\ell^{2}$-norm sampling, we first generate a random number $a$ from the uniform distribution $\mathcal{U}(0,\|M\|_{F}^{2})$ and perform an efficient binary search algorithm to find the leftmost index such that $a$ is less than or equal to the corresponding cumulative sum. Through repeating the $\ell^{2}$-norm sampling for $k$ times, we get the small Q-contexts $R$. We normalize each row to get $\tilde{R}$.

In addition, we can use column sampling to reduce the matrix $\tilde{R}$ again. Applying the theorem and similar analysis to $\tilde{R}^{\top}$, we get a smaller matrix $\tilde{C}$ for which with high probability

Then

We obtain good approximations of the right singular vectors $U$ and singular values $\Sigma$ of $\tilde{R}$, by simply doing the same calculations for the much smaller matrix $\tilde{C}$.

Step 3: Calculating word embedding from Q-contexts. It is known (Mnih and Kavukcuoglu, 2013) that embedding matrix $E_{w}$ taken to be the form $E_{w}=V_{w}\sqrt{\Sigma}$ can be beneficial in predictive performance, where $V_{w}$ and $\Sigma$ are right singular vectors and singular values of information matrix $M$. In view of Theorem (4.2) later, with high probability $\tilde{R}^{\top}\tilde{R}\approx M^{\top}M$, which implies that $\tilde{R}\approx U\Sigma V_{w}^{\top}$ where $U$ is the left singular vectors of $\tilde{R}$. Since $U$ is a (partial-)isometry, we obtain the matrix form of Equation (1):

where $\Lambda=D^{-1}U\Sigma^{-1/2}$ is the matrix of context coefficients.

Complexity Analysis. We get Algorithm 2 by putting the above procedures together. As for line 1, it requires $\mathcal{O}(\text{nnz}(X))$ time to perform point-by-point operations on the nonzero elements of the co-occurrence matrix $X$. As for line 2 and line 3, time complexity for state preparation is $\mathcal{O}(n)$ and for $\ell^{2}$-norm sampling is $\mathcal{O}(k\log n)$. As for line 4, $\mathcal{O}(\text{poly}(k))$ time is spent in singular vectors computation. In MF method, it needs to take highly expensive $\mathcal{O}(\text{poly}(n))$ time to compute singular vectors of original information matrix $M$. As for line 5, $\mathcal{O}(kdn)$ time is spent in matrix multiplication.

We remark that that approximation in Theorem 4.2 is of high quality, both in theory and practice.

From a mathematical point of view, approximation error in Theorem 4.2 is small provided that $k~{}\geq~{}\frac{3\|\varsigma\|_{op}^{2}}{\epsilon^{2}}\log{\text{sr}}(\varsigma)$ rows are sampled from the information matrix $M$. Our analysis is tighter than (Frieze et al., 2004) since we make use of a stronger bound on matrix concentration (Theorem 4.1) concerning stable rank and spectral norm of a matrix, instead of dimension and Frobenius norm. Stable rank is upper bounded by the rank of a matrix, which is more effective in a low-rank matrix regime such as word embedding problem. Spectral norm is bounded above by Frobenius norm as well. Moreover, the Frobenius norm can fail to capture the genuine behavior of random matrices in even very simple examples (Tropp, 2015). Therefore, our analysis shows that with very few samples of rows, $\tilde{R}^{\top}\tilde{R}$ can be a good approximation of $M^{\top}M$ with negligible error.

These theoretical observations are also consistent with our experiments. Indeed, in Section 5.5, we will see that with information matrix dimensions varying from 50k to 400k, the number of row samples required to achieve a good approximation is relatively stable — in our test it ranges from 40k to 50k, while less than only $2\%$ of accuracy is lost compared to direct calculation with the original information matrix.

Training corpora. We conduct our experiments on three large scale English corpora: enwik9 with about 0.1 billion tokens, WebBase (Han et al., 2013) with about 3 billion tokens, and the filtering of Common Crawl corpus (CC) (Ortiz Suárez et al., 2020) with about 6 billion tokens. We first lowercase text and then remove the noisy text like HTML span. We mainly implement data pre-processing on the basis of the script¹¹1http://mattmahoney.net/dc/textdata.html used in word2vec. After pre-processed, the vocabulary sizes of three corpora are 56,466 (on enwik9), 277,704 (on WebBase) and 400,000 (on CC) respectively. We prepare co-occurrence matrices with a context window size of 10. The co-occurrence matrices are square symmetric matrices whose row sizes are equal to vocabulary sizes. Our training corpora (enwik9²²2http://mattmahoney.net/dc/enwik9.zip,WebBase³³3http://ebiquity.umbc.edu/redirect/to/resource/id/351/UMBC-webbase-corpus,CC⁴⁴4https://oscar-public.huma-num.fr/shuff-dedup/en/) are publicly available. Note that the original CC is super-scale corpora, which obtains over 418B tokens. To facilitate experiments, we take 6B tokens. The statistics of the three raw corpora are listed in Table 2.

Baselines. In our experiments, we mainly compare our method WEQ with popular benchmarks listed below.

• •

  MF (Levy and Goldberg, 2014) is a global matrix factorization method that uncovers the semantic information implied in the matrix, such as PPMI and SPPMI matrices.

• •

  word2vec⁵⁵5https://code.google.com/archive/p/word2vec/source/default/source (Mikolov et al., 2013b) is a two-layer neural network that is trained to reconstruct linguistic contexts of words. We choose the SGNS architecture in our experiments.

• •

  GloVe⁶⁶6https://github.com/stanfordnlp/GloVe (Pennington et al., 2014) is a typical co-occurrence count based neural method, which captures global information from word co-occurrence matrix in the training corpus.

• •

  fasttext⁷⁷7https://github.com/facebookresearch/fastText (Bojanowski et al., 2017) is a character-based method which represents each word as an n-gram of characters.

In our experiments, we use code⁸⁸8https://github.com/stanfordnlp/GloVe/blob/master/src/cooccur.c from GloVe to obtain co-occurrence matrix. For MF and WEQ, matrices are stored in a list of lists (LIL) sparse matrix format to allow efficient row operations. The shift parameter of SPPMI is set to ${\kappa}$=5 on all training corpora. Sample size $k$ for our method WEQ are chosen from [10, 100000] as shown in Fig. 2. We factorize the Q-contexts matrix $\tilde{C}$ to get singular vectors and singular values. According to the official guide, we use 15 iterations to train GloVe and word2vec methods and 5 iterations to train fasttext method. The number of dimensions of embedding is set to 300. Note that we set the dimension of word embedding to sample size $k$ when $k$ is smaller than 300 in ablation study of sample size. All experiments are carried out on a cloud server with Intel(R) Xeon(R) Gold 6231C CPU. We set the number of CPU cores to 20, 20, 10, 1 and 1 for fasttext, word2vec, GloVe, MF and WEQ, respectively.

To measure the quality of embeddings, we conduct experiments on three different evaluation tasks: Word Similarity, Text Classification and Named Entity Recognition (NER).

Word similarity. The word similarity task is to measure how well the semantic relationship between words is captured by word embeddings. Experiments are conducted on MEN (Bruni et al., 2014) and WS353 (Agirre et al., 2009) datasets. To estimate the quality of word embeddings, we compute the Spearman Rank Correlation score between the word embedding cosine similarities and human-annotated scores.

Text classification. In the text classification task, the classifier predicts predefined labels for the given texts. We conduct experiments on two datasets, i.e., 5AbstractsGroup (5AG) (Liu et al., 2018) and Stanford Sentiment Treebank (SST-2) (Socher et al., 2013).We choose TextCNN (Kim, 2014) as our classifier, and measure the performance by weighted F1 metrics. The sizes of the kernels are 2, 3 and 4 respectively. We use the Adam optimizer with a mini-batch size of 128 and learning rate of 0.001.

Named entity recognition. NER is the task of identifying and categorizing the entities in the given text. We conduct experiments on two benchmark datasets: Conll (Tjong Kim Sang and De Meulder, 2003) and BTC (Derczynski et al., 2016). We choose wordLSTM+charCNN+CRF (Ma and Hovy, 2016) as our NER model and measure the performance by entity-level F1 metrics. The wordLSTM is built upon a one-layer bidirectional LSTM structure with 200 hidden size. We employ the SGD algorithm to optimize model parameters on mini-batches of size 100, with a learning rate of 0.001.

Table 3 illustrates the results on various evaluation tasks and training times of all word embedding methods. It can be observed that our method WEQ outperforms all baselines in three corpora on training time (and resources) metric while the performance metrics are fairly close.

In terms of evaluation tasks, we notice that MF and WEQ are superior to GloVe, word2vec and fasttext in most word similarity tasks, whereas performing slightly worse on text classification and NER tasks. For both NLP tasks (Text classification and NER), models using pre-trained word embeddings have remarkable improvement compared to the one using random embeddings. Comparing with MF, results obtained with WEQ are fairly close to the ones obtained with MF and sometimes even better, despite the fact we only use a fraction of the original information matrix.

We also observe that the released GloVe embedding is about 3% better than word embeddings trained by us on Conll task. One possible reason is the mismatch of vocabulary (Ma and Hovy, 2016). We reuse the data pre-process script of word2vec excluding punctuations and digits. And the punctuations and digits are important in the Conll task.

To verify the stability of our method WEQ, we run 5 times for each embedding. We only compute the mean and variance of WEQ because other methods need a long time to train word embedding. The small variance shown in Table 3 indicate that the word embedding obtained by WEQ are stable.

With regard to training time, MF method innately has an advantage over neural network based methods (word2vec, GloVe and fasttext) given the same computational settings. Our method WEQ improves significantly on top of that. As can be seen from the Table 3, our method WEQ-PPMI only requires 54.38% (on enwik9), 15.32% (on WebBase) and 7.72% (on CC) of the training time of MF-PPMI respectively. With the CC corpus, while GloVe needs 270.5 minutes with 10 CPU cores to obtain word embeddings, it takes only 31.4 minutes with 1 CPU core with WEQ-PPMI. Although, as the vocabulary size increases, the training time of WEQ increases accordingly, the time taken by MF increases by a significantly larger margin.

Note that in the case closest to practice (6B tokens and 400K vocabularies), we achieve at least $11\sim 13$ times speed-up compared with other baselines. In more realistic scenarios where the number of tokens is trillions and the size of vocabulary is several million, one can infer from the growing trend that the speed advantage of our algorithm could potentially reach several orders of magnitude.

The analysis presented above confirms that WEQ method can substantially minimize the training time while maintaining benchmark performance.

Time analysis. Recall three main steps of the WEQ method: information matrix computation, Q-contexts construction (including state preparation and sampling) and calculating word embedding from Q-contexts (including computation of singular values and left singular vectors, and embedding computation). The breakdown of computational time is displayed in Table.4. Note that the Q-contexts construction step takes up only no more than 11% of the whole time, which shows the efficiency of Q-contexts construction step.

Sample size. In our method, Q-contexts matrix is obtained by sampling from the information matrix. To investigate how the sample size affects the performance of WEQ, we carry out experiments with different sample sizes on all three corpora as shown in Fig.2. We present a representative benchmark score, similarity task MEN, against the training times under corresponding sample sizes. The results of other evaluation tasks are shown in Fig.3. As the sample size increases, it is obvious that the performance of WEQ on MEN approaches that of MF. The convergence is fairly fast, especially when the corpus size is large. We can see that to reach at least 98% of the performance of MF, the sample size of WEQ method only needs to be 40k (on enwik9), 50k (on WebBase), 50k (on CC), which are about 70.83%, 18.00%, 12.50% of the original size of the information matrices, respectively. This validates the theoretical observations in Section 4.1. As the size of the original information matrix increases, the percentage of samples needed decreases. In terms of training time, it rises slowly at small sample sizes then trends upward sharply at larger sample sizes. This indicates that the training time can be effectively reduced when sampling a small matrix to get word embedding.

$\ell^{2}$-norm sampling vs uniform sampling. In WEQ method, we use $\ell^{2}$-norm sampling to capture the typical information of origin information matrix. To verify the effectiveness of the $\ell^{2}$-norm sampling, we replace the $\ell^{2}$-norm sampling with uniform sampling in the Q-contexts construction step. Uniform sampling means that each row of information matrix $M$ will be sampled with equal probability. We compare the performance of word embeddings trained by these two different sampling algorithms. The evaluation results are shown in Fig. 4. $\ell^{2}$-norm sampling has great advantages over uniform sampling in terms of word similarity scores. In terms of downstream tasks (text classification, NER), $\ell^{2}$-norm sampling has $1.5\%\sim 4\%$ increase compared with uniform sampling. These experimental results are consistent with our theoretical proof in Section 4.

Nonzeros/Sparsity analysis. To investigate the correlation between matrix sparsity and algorithm efficiency, we show a comparison between time and number of nonzeros in an information matrix $M$ in Fig. 5. By TES, we mean total time excluding SVD computation, which corresponds to Algorithm 3 excluding line 3. Since the time for SVD computation depends on the SVD engine used in our method, we do not provide a detailed analysis. The Pearson Correlation Score between $\text{nnz}(M)$ and TES is 0.9583. It demonstrates that our algorithm efficiency is approximately linearly and positively correlated with the number of nonzero elements of $M$. We also observe that number of nonzero elements in SPPMI matrix is significantly less than that in PPMI. That is because only the positive values higher than $\kappa$ are retained after the shifted operation. The decrease of number of nonzero elements leads to a more sparse matrix and shorter TES.

Constructing the word embedding which could express the semantic features is a fundamental task in Natural Language Processing. It brings benefits to many NLP tasks (Maruf and Haffari, 2018; Khabiri et al., 2019). Briefly, recent work about word embedding can be categorized into two genres, i.e., neural network based methods (Mikolov et al., 2013b; Bojanowski et al., 2017; Pennington et al., 2014) and global matrix factorization based methods (Deerwester et al., 1990; Levy and Goldberg, 2014; Arora et al., 2016). It is found that we can bridge these two categories by unifying neural network based methods into a matrix factorization framework (Arora et al., 2016; Levy and Goldberg, 2014).

Most of the neural network based methods model the relation between the target word and its contextual words. Among them, the SGNS (Skip-Gram with Negative Sampling) model is the popular implementation of word2vec (Mikolov et al., 2013b). It separates local context windows on the whole corpus, and focuses on maximizing the likelihood of contextual words based on the given word. Fasttext (Bojanowski et al., 2017) is an extension of the word2vec model, which represents each word as an n-gram of characters by a sliding window. GloVe (Pennington et al., 2014) is trained based on the global word-word co-occurrence counts from a corpus. As for one kind of global matrix factorization based method, MF (Levy and Goldberg, 2014) factorizes the SPPMI matrix explicitly, rather than iteratively tuning the network parameters.

More recently, ELMo (Peters et al., 2018) and BERT (Devlin et al., 2019) achieve state-of-the-art results in a variety of NLP tasks. Despite the state-of-the-art results achieved by deep contextualization models, pre-trained word embedding should not be neglected. Indeed, ELMo is dependent on GloVe embedding as input during training (Peters et al., 2018). Moreover, both ELMo and BERT are extremely large deep networks that require huge computing resources. The pre-training of $\text{BERT}_{large}$ takes 4 days to complete on 64 TPUs (Devlin et al., 2019). And it is also difficult to apply deep contextualization models directly on low-resource PCs or mobile phones (Strubell et al., 2019; Sun et al., 2020).

In the studies described above, researchers paid less attention to the computational cost of word embedding training. It is true that large-scale word embedding training requires significant computational costs (Mikolov et al., 2013b; Mnih and Kavukcuoglu, 2013; Pennington et al., 2014; Mikolov et al., 2013a). In this work, we aim to present a novel efficient method to obtain word embeddings of a large-scale vocabulary, while maintaining its superiority in terms of effectiveness.

Low-rank approximation problems in mathematical modeling have been studied for decades (Frieze et al., 2004; Dahiya et al., 2018), which is approximating a matrix by one whose rank is less than that of the original matrix. The aim is to obtain more compact information representation and less complex data modeling with limited losses. It arises in many applications such as recommendation system (Arrazola et al., 2020; Tang, 2019). Unlike their implementations of random matrices and matrices of small sizes (Arrazola et al., 2020), our WEQ investigates the actual large matrices associated with natural language and reveals new relations for word embeddings. Our work is also different from (Tang, 2019; Frieze et al., 2004) since we provide more detailed theoretical and empirical analysis on matrix concentration, and direct analysis of singular vectors. We also have novel treatments for sparse matrices and better algorithmic designs for state preparation.