How to speed up Word Mover's Distance computation on text in dataframe?

Viewed 236

I have a dataframe that looks like but it is around 1000 rows:

 ID    Messages   
 1     I eat apples
 2     I eat oranges
 3     I like to ski
 4     I do not like vegetables. 
 5     The sky is blue.  

I want to calculate the word mover's distance between all messages. My code currently looks like:

 import gensim.downloader as api
 #importing word vectors 
 wiki_vectors = api.load('glove-wiki-gigaword-50')

 #initializing word2vec model
 base_model = Word2Vec(vector_size = 50, min_count= 5)
 base_model.build_vocab(wiki_vectors)

 #using model to get wordmover's distance
 distance_results =  pd.DataFrame([[base_model.wv.wmdistance(p1, p2) for p2 in data.Messages.str.split()] 
                     for p1 in data.Messages.str.split()]
                   , columns = data.ID
                   , index = data.ID )

However, this takes a very long time to run and won't be feasible on larger datasets. I have tried using modin.pandas but that didn't change anything. I have looked into multiprocessing but I need to get the results in the dataframe.

I thought maybe use swifter and create a function to use apply with but I can't seem to get it right.

 def compute_wmd(col):
       return [base_model.wv.wmdistance(p1, p2) for p2 in data[col].str.split() for p1 in data[col].str.split()]


 data['wmd'] = data['Messages'].swifter.apply(compute_wmd)

Any help would be greatly appreciated!

2 Answers

In my opinion you are asking a question at the frontier of available tools. Whenever I reach that I think through the bottlenecks of the codebase in order to develop a potentially more efficient approach.

Word Mover's Metric is symmetric (i.e. agnostic to the ordering of arguments), so you can reduce you calculation to only calculating the lower triangular portion and ignoring the diagonal (same input args). This should reduce your runtime by more than 50% as you are iterating through serially.


partial_results = dict()

# m1, m2 below can be mapped to your data.ID values to be consistent
for m1, p1 in enumerate(data.Messages.str.split()):
    for m2, p2 in enumerate(data.Messages.str.split()):
        if m1 <= m2: 
            continue
        partial_results[(m1,m2)] = base_model.wv.wmdistance(p1, p2)

df = pd.DataFrame.from_dict(partial_results, orient='index')

df here will be a normalized form for your dataset, instead of being a proper matrix. You can then recreate your matrix if needed by swapping m1, m2 index values and appending records to the dataframe.

Other speedups or vectorizations may be possible, but the above approach is amenable to multiprocessing (i.e. is "embarrassingly parallel").

Other items to consider to make more efficient

  • refactoring and profiling the word mover source code
  • reducing the model size (glove-wiki-gigaword-50 is large, 6B words to compare)
  • invoking numba on a refactored source code

Note that one aspect of your current code is nonsensical:

#initializing word2vec model
base_model = Word2Vec(vector_size = 50, min_count= 5)
base_model.build_vocab(wiki_vectors)

There's no need to create a whole Word2Vec model instance – that's only needed if you were going to train your own word vectors from your own training texts. If you're just interested in loading some other pretrained vectors – as here glove-wiki-gigaword-50 into the variable wiki_vectors – you can just use that instance of wiki_vectors to compute your wmdistance() values.

Still, once you've got a set fo word-vectors, doing the pairwise WMD calculations is an inherently expensive calculation. Rather than the simple set of single-vector to single-vector computations that are used in, say, cosine-similarity, the WMD value is the result of a more sophisticated optimization-search, between what are essentially two "bags of word-vectors". This pairwise calculation especially grows in complexity as each text grows in word-count - so if manageable between a bunch of 5-word phrases, may become impracticallly time-consuming between 20-word or 100-word texts.

Some ways to make a large batch of WMD calculations go faster without reaching into the source code for Gensim's wmdistance() method could include:

  • As ~tom-roderick mentions in the sibling answer, make sure you're not redundantly calculating both WMD(A,B) and WMD(B,A) - as they'll be identical.

  • Pre-slimming your texts to eliminate less-important words. You'd then be making a faster calculation between smaller sets of words – which isn't the same as the full WMD, but might be as good or even better, depending on your application & how well your heuristic for discarding some words works. What words might be less important? Perhaps: the most-frequent X% of words.

  • Since in many cases, you'd be interested in the top-N most-similar texts for any other text – rather than all the pairwise WMD values – you may be able to simply avoid doing many of the calculations, if some pre-comparison suggests they are unlikely to be very close. For example, WMD is usually a better measure of similarity than naively summarizing each text as the average of all its word-vectors, then using cosine-similarity. But that naive calculation is much faster, and any texts that aren't in, say, the top half of all most-similar texts by that (fast) calc almost certainly won't be in the top-N (for small N) of closest-by-WMD.

Separately, if you wanted to change the wmdistance() code, some other avenues for optimization would avoid redundant calculations when doing large batches of texts, and might offer a practical speedup. For example:

  • Currently, for each call with a pair of texts, that method re-builds a table of all their local words' pairwise similarities before doing the WMD optimization-search. Conceivably, in large batches with many of the same words repeated, that coudl be done far fewer times.

  • If looking for the top-N least-distant texts, there might be cases where, given your current top-N candidates, a current calculation discovers that the distance must be larger than your current best, long before the exact distance is calculated. The method could conceivably short-circuit quick-return in such a situation, with just an indication: "farther than you care about".

If your need for speed is intense enough, and you'll be repeating this often enough over many production cycles, you might want to look into making or contracting for those kinds of optimizations into the existing Gensim wmdistance() support. But if you only need to do this occasionally, or in a project that can suffer a few waits, or perhaps surge-hire a large batch of machines to each do a small portion of the work separately then merge results, then you'd just want to use the current code, with its inefficiencies, as best you can.

Related