Skip-Gram with Negative Sampling

The problem, from first principles

Word embeddings are learned from a simple game: given a word, predict the words around it. In the skip-gram model each word has two vectors, one used when it is the center and one when it appears as context. Scoring every possible context word with a softmax over the whole vocabulary is far too expensive, so negative sampling turns the problem into binary classification: the true (center, context) pair should score high, and a handful of randomly drawn negative words should score low. The score of a pair is the dot product of the two vectors, passed through a sigmoid. Training pulls vectors of words that share contexts together, which is why similar words end up with similar vectors.

From theory to code

Implement skipgram_pairs and sgns_loss.

Constraints

• skipgram_pairs(token_ids, window) returns a list of (center, context) pairs. For each position i in order, the contexts are the tokens at positions i - window .. i + window (inside the sequence, excluding i), in increasing position order.
• W_in and W_out are (V, d) arrays: W_in[c] is the center vector and W_out[o] the context vector. sgns_loss(W_in, W_out, center, context, negatives) takes integer arrays center and context of shape (B,) and negatives of shape (B, K).
• For each example the loss is -log sigmoid(u_o . v_c) - sum_k log sigmoid(-u_k . v_c). Return the mean over examples as a float. Use np.logaddexp(0, -x) for -log sigmoid(x).

Hints