Negative Sampling › Comparing Full Softmax and Negative Sampling

class SkipGramFullSoftmax:
    """Skip-gram with full softmax for comparison."""

    def __init__(self, vocab_size, embedding_dim):
        self.W = np.random.randn(vocab_size, embedding_dim) * 0.01
        self.W_prime = np.random.randn(vocab_size, embedding_dim) * 0.01
        self.vocab_size = vocab_size
        self.embedding_dim = embedding_dim

    def forward(self, center_idx, context_idx):
        center_vec = self.W[center_idx]
        scores = self.W_prime @ center_vec
        exp_scores = np.exp(scores - np.max(scores))
        probs = exp_scores / exp_scores.sum()
        loss = -np.log(probs[context_idx] + 1e-10)
        return loss, probs, center_vec

    def backward(
        self, center_idx, context_idx, probs, center_vec, learning_rate
    ):
        grad_probs = probs.copy()
        grad_probs[context_idx] -= 1
        for j in range(self.vocab_size):
            self.W_prime[j] -= learning_rate * grad_probs[j] * center_vec
        grad_center = self.W_prime.T @ grad_probs
        self.W[center_idx] -= learning_rate * grad_center

    def train_pair(self, center_idx, context_idx, learning_rate=0.025):
        loss, probs, center_vec = self.forward(center_idx, context_idx)
        self.backward(center_idx, context_idx, probs, center_vec, learning_rate)
        return loss

    def get_embedding(self, word_idx):
        return self.W[word_idx]


# Train both models with the same random seed for fair comparison
np.random.seed(42)
model_softmax = SkipGramFullSoftmax(vocab_size=len(vocab), embedding_dim=30)
model_negsamp = SkipGramNegSampling(
    vocab_size=len(vocab),
    embedding_dim=30,
    word_counts=word_count_by_idx,
    num_negatives=5,
)

losses_softmax = []
losses_negsamp = []

for epoch in range(50):
    np.random.shuffle(training_pairs)
    epoch_loss_sm = 0
    epoch_loss_ns = 0

    for center_idx, context_idx in training_pairs:
        epoch_loss_sm += model_softmax.train_pair(
            center_idx, context_idx, learning_rate=0.1
        )
        epoch_loss_ns += model_negsamp.train_pair(
            center_idx, context_idx, learning_rate=0.1
        )

    losses_softmax.append(epoch_loss_sm / len(training_pairs))
    losses_negsamp.append(epoch_loss_ns / len(training_pairs))