Learning Word Embedding with Better Distance Weighting and Window Size Scheduling

In the CBOW model, on a sliding window of the text segment $(w_{t-r},...,w_{t-2},w_{t-1},\\ w_{t},w_{t+1},w_{t+2},...,w_{t+r})$ with window size $r$, we predict the center word $w_{t}$’s occurrence probability using the context words $C_{t}=(w_{t-r},...,w_{t-1},w_{t+1},...,w_{t+r})$ as

Here, $E$ is the neural network weight matrix of size $d\times|V|$, and $x_{t}$ is the one hot representation of word $w_{t}$ with element 1 at position $t$ and zero on all other positions. So, $u_{w_{t}}=Ex_{w_{t}}$ picks the $t$-th column of the $E$ matrix, and this column therefore can be viewed as a distributed representation of the word $w_{t}$, i.e., the $d$-dimensional embedding of $w_{t}$.

The above model is not distance-sensitive, since the distance between a context word and the target (center) word is absent in Eq.1. However, Our experience says that $w_{t+1}$’s relation with $w_{t}$ is more significant than $w_{t+2}$’s relation with $w_{t}$. Consequently, $w_{t+1}$’s predictive power with respect to $w_{t}$ is generally larger than $w_{t+2}$’s predictive power w.r.t $w_{t}$.

For this reason, it is reasonable to add a distance-related weight $\lambda_{i}$ for each context word $u_{t+i}$ to reflect their influence on predicting $w_{t}$ when averaging them up for $u_{C_{t}}$. Specifically

where $Z=\sum_{|i|\leq r,i\neq 0}\lambda_{i}$ is the normalization factor.

The main question to be answered for this strategy is how the distance-related weights are constructed. We propose the Learnable Formulated Weights (LFW) method to address this issue, which constructs a formula with a small number of learnable parameters that takes distance as input to calculate weights at different distances. This method combines the prior form of the formula with the posterior learning results of the parameters, allowing us to directly model the relationship between weights and distances and perform dynamic adjustments on specific values of parameters.

There are two major aspects of the prior relationship between weights and distances we will model in this paper, each corresponding to two assumptions. The first aspect is whether or not the two context words on both sides of $w_{t}$ with the same distance share the same weight, and the second aspect is which kind of mathematical relationship (exponential or power) better models the relationship between weights and distances on a single side. By combining assumption choices for each of the aspects, we will be able to model the prior relationship in four different ways and find out the best combination for CBOW.

Based on the content above, the formulas corresponding to the third and fourth assumption combinations can be given accordingly as Eq.5 and Eq.6, except that the weights decrease exponentially with respect to distance:

In the experimental part, we will use the above four formulas for weight calculation respectively to demonstrate which of those formulas is the most effective, and at the same time prove the superiority of our method over other approaches to adding weights to CBOW.

For the Skip-gram model, each center word $w_{t}$ is used as an input to predict context words in a sliding window of size $r$. That is, for each context word $w_{t+i}$ with $0<|i|\leq r$, we predict its occurrence probability as

Because no average of context words appears in this model, we cannot make the model distance-sensitive by adding weights for context words. However, since the context words that appear more in the prediction have more influence on how the word embeddings are trained, we can make the model sensitive to distance by sampling more from context words that are closer to the center word and sampling less otherwise. In other words, the probability of being sampled will act in a similar way as the weights added to the CBOW model.

A naive way to do this is to dynamically select a random context window size $r^{\prime}$ from 1 to $r$ for each center word $w_{t}$, and use $w_{t}$ to predict only context words within distance $r^{\prime}$. In this case, the probability of each $w_{t+i}$ with $0<|i|\leq r$ to be used for prediction is

We can learn from Eq.8 that the probability of each $w_{t+i}$ to be used for prediction decreases linearly as the distance $|i|$ increases. Therefore, the context words closer to $w_{t}$ will obtain larger weights than those farther away.

One major problem concerning this dynamic window size strategy is the frequent and irregular changes in window size for each center word resulting from random selections, which ruins the training balance for each word and reduces modeling performance. For example, imagine that the window sizes selected for two center words $w_{t_{1}}$ and $w_{t_{2}}$ are 1 and r respectively, then the number of predictions made by $w_{t_{2}}$ is $r$ times larger than that of $w_{t_{2}}$, causing the model to pay special attention to the center words in some parts while ignoring the training of others, which will lead to an unreasonable allocation of calculations and a decrease in model performance.

We propose the Epoch-based Dynamic Window Size (EDWS) to solve this problem, which cancels the random selection while preserving the dynamicity by gradually increasing the window size by the number of training epochs. In our specific implementation, we use three different window sizes with a ratio of $1:2:3$ to represent the beginning, middle, and end of the training process. The smallest window size appears first, then the middle and largest one, each running for the same number of epochs. Therefore, the window size $r^{\prime}_{k}$ on the $k$th epoch is given by

The EDWS strategy has two major advantages over the original dynamic window size. First, EDWS eliminates the irregularity brought by random selections at its source. Second, EDWS specifies the window sizes to appear in ascending order, forming a gradual learning process from nearby to distant contexts, which aligns with general learning principles. Given these advantages, EDWS is theoretically more effective than the original dynamic window size, which will be demonstrated through experiments in later parts of this paper.

We use two corpora constructed from English Wikipedia for training: enwik9 and text8 ¹¹1Datasets and text preprocessing script are from http://mattmahoney.net/dc/textdata.html. The enwik9 dataset contains the first $10^{9}$ bytes of English Wikipedia with about 120 million words after being processed by Matt Mahoney’s text preprocessing script, and text8 contains the first $10^{8}$ bytes of the preprocessed enwik9 with about 17 million words. We discard words that appear no more than 5 times in each dataset to focus on more valuable words while reducing computation overhead, forming two vocabularies with sizes of 194,377 and 63,643. The enwik9 dataset is much larger than text8, so we conduct pre-experiments on text8 to determine the hyperparameter choice, i.e., the window size. Then, we conduct comparative experiments on enwik9 based on the chosen hyperparameter to demonstrate the effectiveness of our proposed methods.

As for test results, We use the analogical reasoning task dataset proposed in [12] to measure the quality of the trained word embeddings. By evaluating the embeddings’ analogical inference ability, this task reflects the embeddings’ mastery of language and forces a more precise distribution of the trained word embeddings in the vector space to meet the requirements of linguistic relations. There are 8,869 semantic and 10,675 syntactic questions in the test set in total. The semantic questions are divided into five categories, while the syntactic questions are divided into nine categories, as shown in Table 4. An example of these questions is finding the word with the same relative relationship with the word “woman” as the word “man” has with the word “king”. The question should be answered by finding the word embedding that is closest to $vector(``king")-vector(``man")+vector(``woman")$ in cosine distance. The quality of word embeddings is then measured by the total accuracy of all questions in this task.

We list the results of CBOW and Skip-gram with different window sizes in Table 1. Since the dynamic window size strategy is proposed as an original design of Skip-gram, we will use this strategy for all original Skip-gram models used in the experiments, and the window size represents its maximum window size in random selection. We set the word vector dimension to a relatively small value of 128 in this step to reduce the calculation cost of pre-experiments. Consistent with the results reported in [4], we find that the improvement in model performance is significantly related to the increase in window size before the window size reaches 15. After that, the correlation becomes less significant. Considering the impact of increasing the window size on the calculation amount, we choose the window size to be 15 for formal experiments.

In this step, we conduct experiments to show our methods’ effectiveness for Word2Vec on the enwik9 dataset. We set the word vector dimension to 600, following the configuration in [16]. Table 2 and Table 3 show positive experimental results for both proposed methods, where the fonts in bold indicate the highest accuracy for each column. The LFW methods following Eq.3, Eq.4, Eq.5, and Eq.6 improve the overall accuracy of CBOW by 15.3%, 14.4%, 13.3%, and 12.1%, respectively, and EDWS improves that of Skip-gram by more than 2.5%.

Results suggest that formula Eq.3 models the relationship between weight and distance best among all proposed formulas. This may be because: 1) There is no significant difference in the relationship between context on two sides and the center word and using the same parameters on both sides can reduce overfitting. 2) The closest context words have a significantly higher influence on the center word than the other, while the influence of the other words decreases relatively slowly with distance, which can be modeled better with the power function that decreases greatly from distance 1 to 2 and gently at farther distances. To establish an intuitive impression of the relationship between weights and distances obtained using power and exponential functions, we provide Figure 2, which shows the curves formed by the relationship between weights and distances learned from the enwik9 dataset when using Eq.3 (power law decay) and Eq.5 (exponential decay). The curves in the figure take their points at integer distances within the window size, and the sum of the weights for each curve is normalized to 1 for ease of comparison. The behavior of those curves solidifies our analysis of the success of Eq.3. Therefore, we suggest using Eq.3 as the formula for LFW and future context modeling efforts. to achieve the largest performance improvement.

Under their best version, Our methods’ performance exceeds the state-of-the-art accuracy improvement of 13.5% on CBOW and 1.8% on Skip-gram achieved by independent learnable weights presented in [16]. Additionally, the LFW method consumes less than 15% of extra training time, which is a huge improvement compared with the 72% (computed from its 42% slower speed) consumed by the method in [16]. The test results and the improvement amounts of LFW Eq.3 (best version) detailed to each category are shown in Table 4 to illustrate its effectiveness better. Significant relative improvements are brought to nearly every category by our proposed method. The experimental results prove the effectiveness of both our methods.