Word2Vec how to choose the embedding size parameter

Question

I'm running word2vec over collection of documents. I understand that the size of the model is the number of dimensions of the vector space that the word is embedded into. And that different dimensions are somewhat related to different, independent "concepts" that a word could be grouped into. But beyond this I can't find any decent heuristics for how exactly to pick the number. There's some discussion here about the vocabulary size: https://stackoverflow.com/questions/45444964/python-what-is-the-size-parameter-in-gensim-word2vec-model-class However, I suspect that vocabulary size is not most important, but more important is how many sample documents you have and how long they are. Surely each "dimension" should have sufficient examples to be learnt?

I have a collection of 200 000 documents, averaging about 20 pages in length each, covering a vocabulary of most of the English language. I'm using the word2vec embedding as a basis for finding distances between sentences and the documents. I'm using Gensim, if it matters. I'm using a size of 240. Is this reasonable? Are there any studies on what heuristics to use to choose the size parameter? Thanks.

Answer 1

You might find this paper might be the closest thing to what you are looking for if you don't want to treat it as a regular hyperparameter: Towards Lower Bounds on Number of Dimensions for Word Embeddings

The paper claims that there is a lower bound on the embedding based on the corpus. It also purposes a method for finding said lower bound which I will leave the paper to explain since I think I will not do it justice. Here is the most relevant section of the conclusion of the paper:

We discussed the importance of deciding the number of dimensions for word embedding training by looking at the corpus. We motivated the idea using abstract examples and gave an algorithm for finding the lower bound. Our experiments showed that performance of word embeddings is poor, until the lower bound is reached. Thereafter, it stabilizes. Therefore, such bounds should be used to decide the number of dimensions, instead of trial and error.

It has sourced and cited previous work regarding embedding dimension which also might be of interest to you. Unfortunately the conclusion seems to be the following:

As is evident from the above discussion, the analysis of the number of dimensions have not received enough attention. This paper is a contribution towards that direction.

Answer 2

I have checked four well-cited papers related to word embedding: 2013 Word2Vec, 2014 GloVe, 2018 BERT, and 2018 ELMo. Only GloVe has experimented on the embedding dimension for the analogy task (answering "a" is to "b" as "c" is to ?).

Other papers did not report an experiment on embedding dimension size. They are all using an arbitrary dimension on the order of hundreds (100 and 300 are used more frequently). The lack of experiments for embedding size implies that the performance is not very sensitive to this parameter and only the order of magnitude matters, and also other aspects of the model architecture are more important to investigate.

You can carry out a similar experiment using values from three orders of magnitude 10, 100, and 1000. The required dimension size will depend on the task too. I speculate that a classification (discriminative) task would require fewer dimensions than a sentence generation task.

Generally, for hyper parameter optimization, methods like Bayesian Optimization can be used to find the best hyper parameter (here, embedding dimension) with as few (costly) training-evaluations as possible. These techniques are more useful in competitive settings where every improvement counts.

Answer 3

Generally, the exact number of embedding dimensions does not affect task performance.

The number of dimensions can affect training time.

A common heuristic is to pick a power of 2 to speed up training time. Powers of 2 have a good chance to increase cache utilization during data movement, thus reducing bottlenecks.

The most common powers of 2 for word embeddings are 128 or 256, depending on which order of magnitude is preferred.