I feel like I'm missing something obvious here because I can't find any discussion of this. I want to do a lot of reverse lookups (nearest neighbor distance searches) on the GloVe embeddings for a word generation network. I'm currently just iterating through the vocabulary on the cpu. I've sped it up a bit using a process pool, as shown in the snippet below, but it's still very slow for large vocabs.

Is there a way to move this to to the GPU using cuda? I've also read that there is a way to turn this sort of thing into one big matrix operation... Any references would be appreciated. Thanks!

glove = torchtext.vocab.GloVe(name='6B', dim=wordDim)

def closest(vec):
    dists = [(w, torch.dist(vec, glove.vectors[glove.stoi[w]] for w in glove.itos]
    return sorted(dists, key=lambda t: t[1])[0]

output = # word vectors…
# using a process pool to parallelize the lookup
predictedWords = [w for w in list(pool.map(closestWord, output)]

It's usually not very efficient to approach these types of problems in pythonic ways, with list comprehensions and such. This whole process can be done with some matrix math, which will be substantially faster (and able to be computed on the GPU using PyTorch). Using torch.dist with default p=2 will compute the Euclidean distance between two tensors, which is defined as

$$ \text{euclidean distance} = d(\mathbf{v}, \mathbf{w}) = \sqrt{\sum_d \left(\mathbf{v}_d - \mathbf{w}_d\right)^2} $$

You can do this efficiently in PyTorch for every word in your vocab by broadcasting your query word over the whole matrix of word vectors:

def closest(vec):
    dists = torch.sqrt(((glove.vectors - vec) ** 2).sum(dim=1))
    return dists.argmin() # or glove.itos[dists.argmin()] if you want a string output

However, usually people use the cosine similarity to find the closest word vectors rather than the Euclidean distance. This is the cosine of the angle between the vectors, and can be calculated as

$$ \text{Cosine similarity} = \frac{\mathbf{v}\cdot \mathbf{w}}{\|\mathbf{v}\| \|\mathbf{w}\|} $$

which can similarly be implemented in PyTorch as the following:

glove_lengths = torch.sqrt((glove.vectors ** 2).sum(dim=1))

def closest_cosine(vec):
    numerator = (glove.vectors * vec).sum(dim=1)
    denominator = glove_lengths * torch.sqrt((vec ** 2).sum())
    similarities = numerator / denominator
    return glove.itos[similarities.argmax()]

Note that we moved the $\|\mathbf{v}\|$ part of the equation outside the function here, so we don't have to recompute the lengths of each glove vector every time we want to run it.

  • $\begingroup$ Thanks you! I came here this morning to answer my own question after it finally dawned on my that I could just do the element-wise operations on matrices. But your code is more concise than what I had! For the record: closest() as above is 50x faster on the CPU alone than the dumb iteration and it's the same number of lines of code. I don't understand why the GloVe examples don't show it this way. $\endgroup$ Dec 7 '18 at 15:05
  • $\begingroup$ I guess for closest() we could also get rid of the sqrt(). $\endgroup$ Dec 7 '18 at 15:38

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.