Faster Training of Word Embeddings

1 INTRODUCTION . Word embeddings have a long history ( Rumelhart et al. , 1986 ; Bengio et al. , 2003 ; Collobert & Weston , 2008 ) , but have received much attention in recent years due to word2vec ( Mikolov et al. , 2013 ) and its computationally efficient implementation via skip-gram with negative sampling . Word embeddings capture contextual relationships between the words , and have become a standard input representation for the majority of NLP tasks , benefitting , e.g. , classification ( Joulin et al. , 2016 ; Deriu et al. , 2017 ) or machine translation ( Jansen , 2017 ; Conneau et al. , 2017 ) . More recently , state-of-the-art results on many language understanding tasks were achieved by deep transformer architectures such as BERT ( Devlin et al. , 2019 ) , which however are very compute intensive both at training and inference time , even with pre-trained models and reduced parameter space . Thus , simpler and more lightweight static word embeddings such as fastText ( Bojanowski et al. , 2017 ) are still widely used , due to their fast execution , comparable results for particular tasks ( Tseng et al. , 2019 ) , and ability to produce a single vector per word , which helps in information retrieval with interpretability and search index construction . Contributions . In this paper , we present algorithmic and code optimization techniques to improve the training time of word2vec and fastText embeddings on modern general-purpose multicore and manycore computers . We present an optimized open-source implementation of word2vec and fastText that encapsulates a number of algorithmic variants including negative sample sharing , batched updates , and subword units based on byte-pair encoding approach . Our extensive evaluation on three languages shows that the best combinations of optimizations speed up training time by 2.7– 20.6 times while maintaining accuracy of selected NLP tasks . 2 WORD EMBEDDINGS . Word2vec . Word2vec is built upon a simple bilinear regression model trained on word cooccurrence , resulting in numerical feature representations as floating point vectors of dimensionality d. Given a word in a sentence , the goal of the algorithm is to maximize the likelihood of predicting surrounding ( context ) words . To achieve this , the model is trained to increase the probability of predicting particular words if they appear close to a given current word in the training corpus . A popular variant also decreases the probability of predicting words that do not appear close to the current word ( negative sampling ( Mikolov et al. , 2013 ; Goldberg & Levy , 2014 ) ) . During training , the algorithm processes the corpus in a streaming fashion . Each word wi ( called current word ) is processed together with its surrounding context words { wi−C , ... , wi−1 } , { wi+1 , ... , wi+C } , where C is the range of the context window . There are two modes of operation training the model for the following prediction tasks : • Skip-gram ( SG ) : predict target context words using the current word wi as the source . • CBOW : predict the target current word wi using context words as the source . Each word w in the vocabulary of size V is represented as a source ws by one row in the V × d input matrix Min containing word embeddings , and each word is represented as a target wt by one row in the V × d output matrix Mout that is used to calculate the training objective function . The goal of the optimization is to maximize the inner products ( minimize difference ) of real pairs of source current words with the target context words or vice versa , using the binary logistic loss . This approach can be improved by the use of negative sampling , where the algorithm additionally maximizes the difference between the source current words and the words picked randomly out of the source ’ s context . Training is performed using stochastic gradient descent ( SGD ) . SGD is performed in parallel with p threads by splitting the training corpus into p parts and processing them asynchronously ( “ Hogwild ” approach ( Recht et al. , 2011 ) ) . The final embedding of each word is its corresponding row in Min . Mout is discarded at the end of training . FastText . FastText ( Bojanowski et al. , 2017 ) improves word2vec by utilizing subwords of target words during the training . A typical run of fastText uses subwords of lengths k = 3 . . . 6 , containing delimiters 〈 , 〉 at the boundaries . For example , for the word paris and k = 3 the subwords are : 〈pa , par , ari , ris , is〉 . In fastText , the embeddings Min are extended to contain rows representing both entire words as well as hashes of all their subwords . Additionally , the representation of the entire word is added to the set of its subwords . During the execution of the algorithm , the hidden layer h is built by averaging vectors in Min representing the source word ’ s subwords . The final vector embedding for each word is obtained in the same way . Mout remains unchanged . Fig . 1 shows an example of how the word vectors are stored and accessed . A single update is described in Alg . 1 . Related work . FastText has been implemented as a part of the popular Gensim library ( Rehurek & Sojka , 2011 ) using Cython and a default machine ’ s BLAS library ( e.g. , Intel MKL ) for algebraic operations . In our experiments we found the code memory-expensive and slow : training 5 epochs on a 1 GB English Wikipedia dump with 24 threads took approximately 11 hours on a Knights Landing CPU , about 10 times slower than the original fastText . Therefore , we use the original code provided by Facebook Research ( 2016a ) as the baseline in all our experiments . For skip-gram with negative sampling , pWord2Vec ( Ji et al. , 2016 ) transforms the “ Hogwild ” approach into “ Hogbatch ” , by performing updates on multiple context words at once ( effectively turning a series of dot products into a matrix-matrix operation ) and sharing negative samples for the entire batch . We employ similar techniques in our implementation . The work by Rengasamy et al . ( 2017 ) extends this approach by context combining , where multiple contexts can share a set of negative samples and be updated all at once . We do not adapt this approach as it requires careful preprocessing rewarded by a relatively small speedup over pWord2Vec . Algorithm 1 : A single iteration of the original fastText algorithm . In skip-gram , it is performed on each current-context word pair ( as source-target ) . In CBOW , all context words are used as source words at the same time . Data : source word ( s ) ws , target word wt , learning rate l , number of negative samples n 1 if skip-gram then // Initialize . 2 h = Min ( ws ) + ∑ z∈subwords ( ws ) Min ( z ) count ( subwords ( ws ) ) +1 // Average vectors of the source word ( s ) and their 3 else if CBOW then // subwords to obtain the hidden layer . 4 h = ∑ s ( Min ( ws ) + ∑ z∈subwords ( ws ) Min ( z ) ) ∑ s ( count ( subwords ( ws ) ) +1 ) 5 g = 0 // Reset the gradient . // Update the target word . 6 α = l ( 1− σ ( h ·Mout ( wt ) ) ) // Compute positive score reflecting similarity between // h and the row Mout ( wt ) representing wt . 7 g = g + α ·Mout ( wt ) // Build the gradient . 8 Mout ( wt ) =Mout ( wt ) + α · h // Update the target word . 9 for t′ ← 1 to n do // Update negative samples : negative score . 10 pick a random negative sample wt′ 6= wt 11 α = l ( 0− σ ( h ·Mout ( wt′ ) ) ) ) // Compute negative score . 12 g = g + α ·Mout ( wt′ ) ) // Build the gradient . 13 Mout ( wt′ ) ) =Mout ( wt′ ) ) + α · h // Update the target word . 14 end 15 if skip-gram then // Update the source rows ( s ) . 16 Min ( ws ) =Min ( ws ) + g 17 foreach z ∈ subwords ( ws ) do 18 Min ( z ) =Min ( z ) + g 19 end 20 else if CBOW then 21 foreach ws do // As a result , the difference between rows in Min ( ws ) 22 Min ( ws ) =Min ( ws ) + g // and Mout which corresponds to the positive samples 23 foreach z ∈ subwords ( ws ) do // drops , and the difference between rows in Mout 24 Min ( z ) =Min ( z ) + g // which correspond to the negative samples increases . 25 end 26 end Word2vec and fastText have been also implemented for GPU clusters . BlazingText ( Gupta & Khare , 2017 ) tackles the problem of efficient batch size and synchronization for multiple GPUs . While this issue is of no concern on CPU , they report the execution time on a single GPU comparable to a 16- threaded CPU fastText baseline . We further speed up the CPU implementation . The work by Bae & Yi ( 2016 ) reports up to 11× speedup of word2vec with negative sampling run on a K20 GPU over the single-threaded CPU word2vec . However , they report only up to 1.6× speedup over a 12-threaded CPU run . Our no subword version of the code are roughly 5 ( skip-gram ) and 6 ( CBOW ) faster than the 20-threaded runs of the original word2vec . Word2vec and fastText are memory-intensive algorithms . Additionally , fine-grained parallelism is limited by relatively small vectors typically used in the computations . These characteristics severely limit the potential advantages of GPU over CPU . Li et al . ( 2019 ) discuss a distributed version for many GPUs aiming at the reduction of write conflicts in updates . Similarly ( and independently ) , we made attempts at pre-scheduling a list of currentcontext word updates , but we found the overhead of this preprocessing prohibitive . Nonetheless , the algorithmic variants presented in our paper can be applied in distributed setting , as long as the data used for a single update fits inside a batch used in the distributed computation . This will be the case for our variants , since they either execute separate updates on each current-context word pair , or update current words with their entire ( typically small ) contexts . This is also the case in the original fastText and therefore , the communication cost should not increase . Another popular word embedding model is GloVe ( Pennington et al. , 2014 ) . While the Authors claim superiority over word2vec , a more thorough evaluation ( e. g. , by Wang et al . ( 2019 ) , or Kumar et al . ( 2020 ) ) shows that there is no clear winner , as the results may vary depending on the training corpus , evaluation task and hyperparameters used . Additionally , GloVe lacks the information on word morphology provided by fastText , and does not scale well for large vocabularies . Since GloVe is based on a completely different algorithmic structure , that is , creation and reduction of a global word co-occurrence matrix , there is no straightforward way to apply to it our code optimizations and variants .

Summary: This paper applies several algorithmic and code optimization techniques to reduce the training time for word2vec and fastText. The final implementation runs 3-20x times faster than the baselines on manycore CPUs, with almost no quality loss. The improved implementation is a result of the combination of many techniques, including namely no_subword, Minibatching, Negative sharing etc.