Home

Awesome

Swivel in Tensorflow

This is a TensorFlow implementation of the Swivel algorithm for generating word embeddings.

This is the source{d}'s fork, which is different from the original. See "Changes in this fork".

Swivel works as follows:

  1. Compute the co-occurrence statistics from a corpus; that is, determine how often a word c appears the context (e.g., "within ten words") of a focus word f. This results in a sparse co-occurrence matrix whose rows represent the focus words, and whose columns represent the context words. Each cell value is the number of times the focus and context words were observed together.
  2. Re-organize the co-occurrence matrix and chop it into smaller pieces.
  3. Assign a random embedding vector of fixed dimension (say, 300) to each focus word and to each context word.
  4. Iteratively attempt to approximate the pointwise mutual information (PMI) between words with the dot product of the corresponding embedding vectors.

Note that the resulting co-occurrence matrix is very sparse (i.e., contains many zeros) since most words won't have been observed in the context of other words. In the case of very rare words, it seems reasonable to assume that you just haven't sampled enough data to spot their co-occurrence yet. On the other hand, if we've failed to observed to common words co-occuring, it seems likely that they are anti-correlated.

Swivel attempts to capture this intuition by using both the observed and the un-observed co-occurrences to inform the way it iteratively adjusts vectors. Empirically, this seems to lead to better embeddings, especially for rare words.

Contents

This release includes the following programs.

Building Embeddings with Swivel

To build your own word embeddings with Swivel, you'll need the following:

You'll then run prep.py (or fastprep) to prepare the data for Swivel and run swivel.py to create the embeddings. The resulting embeddings will be output into two large text files: one for the row vectors and one for the column vectors. You can use those "as is", or convert them into a binary file using text2bin.py and then use the tools here to experiment with the resulting vectors.

Preparing the data for training

Once you've downloaded the corpus (e.g., to /tmp/wiki.txt), run prep.py to prepare the data for training:

./prep.py --output_dir /tmp/swivel_data --input /tmp/wiki.txt

By default, prep.py will make one pass through the text file to compute a "vocabulary" of the most frequent words, and then a second pass to compute the co-occurrence statistics. The following options allow you to control this behavior:

OptionDescription
--min_count <n>Only include words in the generated vocabulary that appear at least n times.
--max_vocab <n>Admit at most n words into the vocabulary.
--vocab <filename>Use the specified filename as the vocabulary instead of computing it from the corpus. The file should contain one word per line.

The prep.py program is pretty simple. Notably, it does almost no text processing: it does no case translation and simply breaks text into tokens by splitting on spaces. Feel free to experiment with the words function if you'd like to do something more sophisticated.

Unfortunately, prep.py is pretty slow. Also included is fastprep, a C++ equivalent that works much more quickly. Building fastprep.cc is a bit more involved: it requires you to pull and build the Tensorflow source code in order to provide the libraries and headers that it needs. See fastprep.mk for more details.

Training the embeddings

When prep.py completes, it will have produced a directory containing the data that the Swivel trainer needs to run. Train embeddings as follows:

./swivel.py --input_base_path /tmp/swivel_data \
   --output_base_path /tmp/swivel_data

There are a variety of parameters that you can fiddle with to customize the embeddings; some that you may want to experiment with include:

OptionDescription
--embedding_size <dim>The dimensionality of the embeddings that are created. By default, 300 dimensional embeddings are created.
--num_epochs <n>The number of iterations through the data that are performed. By default, 40 epochs are trained.

As mentioned above, access to beefy GPU will dramatically reduce the amount of time it takes Swivel to train embeddings.

When complete, you should find row_embeddings.tsv and col_embedding.tsv in the directory specified by --ouput_base_path. These files are tab-delimited files that contain one embedding per line. Each line contains the token followed by dim floating point numbers.

Exploring and evaluating the embeddings

There are also some simple tools you can to explore the embeddings. These tools work with a simple binary vector format that can be mmap-ed into memory along with a separate vocabulary file. Use text2bin.py to generate these files:

./text2bin.py -o vecs.bin -v vocab.txt /tmp/swivel_data/*_embedding.tsv

You can do some simple exploration using nearest.py:

./nearest.py -v vocab.txt -e vecs.bin
query> dog
dog
dogs
cat
...
query> man woman king
king
queen
princess
...

To evaluate the embeddings using common word similarity and analogy datasets, use eval.mk to retrieve the data sets and build the tools:

make -f eval.mk
./wordsim.py -v vocab.txt -e vecs.bin *.ws.tab
./analogy --vocab vocab.txt --embeddings vecs.bin *.an.tab

The word similarity evaluation compares the embeddings' estimate of "similarity" with human judgement using Spearman's rho as the measure of correlation. (Bigger numbers are better.)

The analogy evaluation tests how well the embeddings can predict analogies like "man is to woman as king is to queen".

Note that eval.mk forces all evaluation data into lower case. From there, both the word similarity and analogy evaluations assume that the eval data and the embeddings use consistent capitalization: if you train embeddings using mixed case and evaluate them using lower case, things won't work well.

Contact

source{d}'s Machine Learning Team: machine-learning@sourced.tech

Changes in this fork