Word2Vec with Negative Sampling

Train word embeddings with skip-gram and negative sampling, the word2vec method.

You get (center, context) pairs of word ids that appeared near each other. The model has two embedding tables, self.center and self.context. For a true pair, the dot product of the centre vector u and the context vector v should be large. For K random "negative" words it should be small:

loss = -mean over the batch of [ log sigmoid(u . v_pos) + sum_k log sigmoid(-u . v_neg_k) ]

Write three things:

• noise_distribution(counts): the distribution to draw negatives from, counts ** 0.75, normalised to sum to one.
• SkipGramNS: an nn.Module whose forward(center, context, negatives) returns the scalar loss above. center and context have shape (B,), negatives has shape (B, K).
• train_skipgram(pairs, counts, dim): trains a new model, drawing fresh negatives for each batch with torch.multinomial, and returns it.

Function signature

def noise_distribution(counts):

class SkipGramNS(torch.nn.Module):
    def __init__(self, vocab_size, dim): ...
    def forward(self, center, context, negatives): ...

def train_skipgram(pairs, counts, dim):

Example

Input

noise_distribution(torch.tensor([1.0, 16.0, 81.0]))

Output

tensor([0.0278, 0.2222, 0.7500])

Reasoning

The powers are 1 ** 0.75 = 1, 16 ** 0.75 = 8 and 81 ** 0.75 = 27. They sum to 36, so the probabilities are 1/36, 8/36 and 27/36. The power 0.75 gives rare words more negatives than their raw frequency would.

Constraints and notes

• Name the tables self.center and self.context. Both are nn.Embedding(vocab_size, dim).
• Use F.logsigmoid. torch.log(torch.sigmoid(x)) gives -inf when x is a large negative number.
• The tests train on a synthetic corpus where words fall into topics. For at least 90% of the words, the nearest other word by cosine similarity of the centre vectors must be in the same topic.