Negative Sampling › Complete Skip-gram Implementation with Negative Sampling

class SkipGramNegSampling:
    """
    Skip-gram model with negative sampling.

    Maintains separate embedding matrices for center words (W) and
    context words (W'), updating both via the negative sampling objective.
    """

    def __init__(
        self,
        vocab_size,
        embedding_dim,
        word_counts,
        num_negatives=5,
        alpha=0.75,
    ):
        # Embedding matrices
        self.W = (
            np.random.randn(vocab_size, embedding_dim) * 0.01
        )  # Center word embeddings
        self.W_prime = (
            np.random.randn(vocab_size, embedding_dim) * 0.01
        )  # Context embeddings

        self.vocab_size = vocab_size
        self.embedding_dim = embedding_dim
        self.num_negatives = num_negatives

        self.sampler = NegativeSampler(word_counts, alpha=alpha)

    def forward(self, center_idx, context_idx, negative_indices):
        """Forward pass: compute loss and cache values for backprop."""
        center_vec = self.W[center_idx]
        context_vec = self.W_prime[context_idx]
        negative_vecs = self.W_prime[negative_indices]

        pos_score = np.dot(context_vec, center_vec)
        pos_sigmoid = sigmoid(pos_score)

        neg_scores = negative_vecs @ center_vec
        neg_sigmoids = sigmoid(-neg_scores)  # Want these close to 1

        loss = -np.log(pos_sigmoid + 1e-10) - np.sum(
            np.log(neg_sigmoids + 1e-10)
        )

        cache = {
            "center_idx": center_idx,
            "context_idx": context_idx,
            "negative_indices": negative_indices,
            "center_vec": center_vec,
            "context_vec": context_vec,
            "negative_vecs": negative_vecs,
            "pos_sigmoid": pos_sigmoid,
            "neg_sigmoids": neg_sigmoids,
        }
        return loss, cache

    def backward(self, cache, learning_rate):
        """Backward pass: compute gradients and update weights."""
        center_idx = cache["center_idx"]
        context_idx = cache["context_idx"]
        negative_indices = cache["negative_indices"]
        center_vec = cache["center_vec"]
        context_vec = cache["context_vec"]
        negative_vecs = cache["negative_vecs"]
        pos_sigmoid = cache["pos_sigmoid"]
        neg_sigmoids = cache["neg_sigmoids"]

        # Gradient for center word: attractive from positive, repulsive from negatives
        grad_center = (pos_sigmoid - 1) * context_vec
        for i, neg_idx in enumerate(negative_indices):
            neg_sigmoid_pos = 1 - neg_sigmoids[i]
            grad_center += neg_sigmoid_pos * negative_vecs[i]

        # Gradient for positive context embedding
        grad_context = (pos_sigmoid - 1) * center_vec

        # Gradients for negative embeddings
        grad_negatives = np.zeros_like(negative_vecs)
        for i in range(len(negative_indices)):
            neg_sigmoid_pos = 1 - neg_sigmoids[i]
            grad_negatives[i] = neg_sigmoid_pos * center_vec

        # Update embeddings (gradient ascent on loss = gradient descent on neg loss)
        self.W[center_idx] -= learning_rate * grad_center
        self.W_prime[context_idx] -= learning_rate * grad_context
        for i, neg_idx in enumerate(negative_indices):
            self.W_prime[neg_idx] -= learning_rate * grad_negatives[i]

    def train_pair(self, center_idx, context_idx, learning_rate=0.025):
        """Train on a single (center, context) pair."""
        negative_indices = self.sampler.sample(
            self.num_negatives, exclude_idx=context_idx
        )
        loss, cache = self.forward(center_idx, context_idx, negative_indices)
        self.backward(cache, learning_rate)
        return loss

    def get_embedding(self, word_idx):
        """Get the embedding vector for a word."""
        return self.W[word_idx]

    def most_similar(self, word_idx, top_n=5):
        """Find most similar words by cosine similarity."""
        word_vec = self.W[word_idx]
        word_vec_norm = word_vec / (np.linalg.norm(word_vec) + 1e-10)
        norms = np.linalg.norm(self.W, axis=1, keepdims=True)
        normalized = self.W / (norms + 1e-10)
        similarities = normalized @ word_vec_norm
        similarities[word_idx] = -np.inf
        top_indices = np.argsort(similarities)[::-1][:top_n]
        return [(idx, similarities[idx]) for idx in top_indices]