Word2Vec Training › Training Word2Vec from Scratch in PyTorch
Language AI HandbookAll chapters Word2Vec Training › Training Word2Vec from Scratch in PyTorchOne section · Sections in this chapter9 / 14 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
While Gensim is the practical choice for most use cases, implementing Word2Vec in PyTorch gives you complete visibility into the training process and makes it easy to modify the architecture. This implementation builds a Skip-gram model with negative sampling, following the mathematical formulation from the previous chapters.
Understanding a from-scratch implementation is valuable not just for academic reasons. When you need to add custom loss terms, modify the sampling strategy, use Word2Vec as a component in a larger architecture, or debug unexpected embedding behavior, the ability to trace the exact computation is essential. Gensim's Cython implementation is fast but opaque. The PyTorch version trades some speed for complete interpretability.
Setup and Data Preparation
We begin by constructing the vocabulary and the negative sampling distribution from the sentences we prepared earlier. The key data structure here is the neg_sample_weights array: this is the unigram distribution that determines how likely each word is to be drawn as a negative sample.
from collections import Counter
# Use the preprocessed data from earlier
# Rebuild vocab from sentences with more flexible min_count
all_tokens = [token for sentence in sentences for token in sentence]
word_counts = Counter(all_tokens)
vocab_words = sorted([w for w, c in word_counts.items() if c >= 1])
word_to_idx_pt = {w: i for i, w in enumerate(vocab_words)}
idx_to_word_pt = {i: w for w, i in word_to_idx_pt.items()}
vocab_size_pt = len(vocab_words)
# Compute word frequencies for negative sampling distribution
total_count = sum(word_counts.values())
word_freqs = np.array([word_counts[w] for w in vocab_words], dtype=np.float32)
# Raise to 3/4 power (Mikolov's unigram^0.75 trick)
neg_sample_weights = word_freqs**0.75
neg_sample_weights /= neg_sample_weights.sum()PyTorch vocabulary size: 91 Top 5 words by negative sampling weight: 'the': freq=12, neg_weight=0.0546 'dog': freq=4, neg_weight=0.0240 'word': freq=4, neg_weight=0.0240 'large': freq=3, neg_weight=0.0193 'and': freq=3, neg_weight=0.0193
The top words by negative sampling weight are typically the most frequent words in the corpus, since the exponent compresses but does not invert the frequency distribution. This is correct behavior: the most frequent words should still appear most often as negatives, just not as disproportionately as their raw frequency would suggest.
Dataset and Negative Sampling
The SkipgramDataset class encapsulates the training pair generation and negative sampling in a single object. PyTorch's Dataset interface requires implementing __len__ (the total number of training examples) and __getitem__ (retrieving a single training example by index). Each item returned is a tuple of (center word index, context word index, negative word indices).
A subtle design choice here is when to sample negatives. We could pre-generate all negatives at the start of each epoch, storing them alongside the training pairs. This would be faster but would mean each training example sees the same negatives every epoch. Instead, we sample negatives on-the-fly in __getitem__. Each time the same positive pair is encountered (in a different epoch or after shuffling), it gets fresh negative samples. This additional stochasticity is beneficial: the model must learn to distinguish the positive pair from many different sets of negatives, which encourages a separation that generalizes across negative samples.
import torch
from torch.utils.data import DataLoader, Dataset
class SkipgramDataset(Dataset):
"""Dataset that generates skip-gram pairs with negative samples."""
def __init__(
self,
sentences,
word_to_idx,
neg_weights,
window_size=5,
num_negatives=5,
):
self.word_to_idx = word_to_idx
self.neg_weights = neg_weights
self.num_negatives = num_negatives
self.vocab_size = len(word_to_idx)
# Generate all (center, context) pairs
self.pairs = []
for sentence in sentences:
indices = [word_to_idx[w] for w in sentence if w in word_to_idx]
for i, center in enumerate(indices):
start = max(0, i - window_size)
end = min(len(indices), i + window_size + 1)
for j in range(start, end):
if i != j:
self.pairs.append((center, indices[j]))
self.pairs = torch.tensor(self.pairs, dtype=torch.long)
def __len__(self):
return len(self.pairs)
def __getitem__(self, idx):
center, context = self.pairs[idx]
# Sample negative words from unigram^0.75 distribution
negatives = np.random.choice(
self.vocab_size,
size=self.num_negatives,
p=self.neg_weights,
replace=True,
)
negatives = torch.tensor(negatives, dtype=torch.long)
return center, context, negatives
dataset = SkipgramDataset(sentences, word_to_idx_pt, neg_sample_weights)
dataloader = DataLoader(dataset, batch_size=128, shuffle=True)Total training pairs: 900 Batches per epoch: 8 Sample training instance: Center word: 'the' (idx=74) Context word: 'quick' (idx=59) Negative words: ['river', 'token', 'jumps', 'encode', 'semantic']
Model Architecture
The Skip-gram model maintains two embedding matrices: one for center words and one for context words. This is the same two-matrix design we derived in the Negative Sampling chapter. To briefly review why two matrices are necessary: during training, a word can appear as both a center word and a context word in different training pairs. If we used a single embedding matrix for both roles, updating a word's embedding as a center word would interfere with its role as a context word. The two-matrix design keeps these two roles separate, allowing independent optimization.
After training, we use the center word embeddings as the final word representations and discard the context embeddings. The center embeddings have received gradient updates from every training pair where the word appeared as the center, which is every occurrence of that word in the corpus (after subsampling). The context embeddings have only received updates from pairs where the word appeared as a context target or negative sample, which is a noisier and less direct signal.
Initialization is also worth understanding. The center embeddings are initialized with small uniform random values in the range where is the embedding dimension. This small-but-nonzero initialization ensures that gradient updates can push embeddings in meaningful directions from the start. The context embeddings are initialized to zero. Zero initialization for context embeddings works because context vectors are immediately updated when they appear as targets or negatives, and the zero initialization provides a neutral starting point that does not bias the first updates.
import torch.nn as nn
class Word2VecSkipgram(nn.Module):
"""Skip-gram model with two separate embedding tables."""
def __init__(self, vocab_size, embedding_dim):
super().__init__()
# Center word embeddings (these are the final embeddings we keep)
self.center_embeddings = nn.Embedding(vocab_size, embedding_dim)
# Context word embeddings (separate matrix, discarded after training)
self.context_embeddings = nn.Embedding(vocab_size, embedding_dim)
# Initialize with small random values
nn.init.uniform_(
self.center_embeddings.weight,
-0.5 / embedding_dim,
0.5 / embedding_dim,
)
nn.init.zeros_(self.context_embeddings.weight)
def forward(self, center, context, negatives):
"""
Compute negative sampling loss for one batch.
Args:
center: (batch_size,) center word indices
context: (batch_size,) positive context word indices
negatives: (batch_size, num_negatives) negative word indices
"""
batch_size = center.shape[0]
# Look up embeddings
center_vec = self.center_embeddings(center) # (B, D)
context_vec = self.context_embeddings(context) # (B, D)
neg_vecs = self.context_embeddings(negatives) # (B, K, D)
# Positive score: dot product of center and context
# Sum over embedding dim, result: (B,)
pos_score = (center_vec * context_vec).sum(dim=1)
pos_loss = -torch.nn.functional.logsigmoid(pos_score).mean()
# Negative scores: center dotted with each negative
# center_vec: (B, 1, D), neg_vecs: (B, K, D)
# bmm gives (B, 1, K), squeeze to (B, K)
neg_score = torch.bmm(neg_vecs, center_vec.unsqueeze(2)).squeeze(2)
neg_loss = -torch.nn.functional.logsigmoid(-neg_score).mean()
return pos_loss + neg_loss
EMBEDDING_DIM = 50
model_pt = Word2VecSkipgram(vocab_size_pt, EMBEDDING_DIM)Model parameters: 9,100 Center embeddings: torch.Size([91, 50]) Context embeddings: torch.Size([91, 50]) Total: 91 words x 50 dims x 2 matrices
The forward pass is worth tracing in detail because it shows how the negative sampling objective is implemented computationally. For a batch of training pairs with negatives each, the computation proceeds as follows. First, we look up the center word vectors (shape ), the positive context vectors (shape ), and the negative context vectors (shape ). The positive loss computes the dot product between each center and its corresponding context, then applies the log-sigmoid and takes the mean. The negative loss uses batch matrix multiplication to compute all center-negative dot products simultaneously, then applies log-sigmoid to the negated scores (because we want the negatives to have low scores) and takes the mean. The two losses are summed and returned.
The torch.nn.functional.logsigmoid function computes in a numerically stable way. This is the correct choice here rather than computing torch.log(torch.sigmoid(x)), which can produce -inf for large negative inputs where sigmoid(x) rounds to zero.
Training Loop with Convergence Monitoring
The training loop implements linear learning rate decay by manually updating the optimizer's learning rate at each step. PyTorch's SGD optimizer does not natively support linear decay, so we compute the current learning rate as a function of the training step and update the parameter group directly.
import torch.optim as optim
def train_word2vec(model, dataloader, epochs=50, lr_start=0.025, lr_min=1e-4):
"""Train Word2Vec with linear learning rate decay and loss tracking."""
optimizer = optim.SGD(model.parameters(), lr=lr_start)
total_steps = len(dataloader) * epochs
step = 0
epoch_losses = []
for epoch in range(epochs):
total_loss = 0.0
for batch in dataloader:
center, context, negatives = batch
# Update learning rate (linear decay)
progress = step / total_steps
lr = lr_start * (1 - progress) + lr_min * progress
for param_group in optimizer.param_groups:
param_group["lr"] = lr
# Forward pass
optimizer.zero_grad()
loss = model(center, context, negatives)
# Backward pass
loss.backward()
optimizer.step()
total_loss += loss.item()
step += 1
avg_loss = total_loss / len(dataloader)
epoch_losses.append(avg_loss)
return epoch_losses
torch.manual_seed(42)
np.random.seed(42)
epoch_losses = train_word2vec(model_pt, dataloader, epochs=100)Training complete. Initial loss: 1.3863 Final loss: 1.3863 Loss reduction: 0.0%
Visualizing Training Convergence
The training loss curve is one of the most informative diagnostics for Word2Vec training. A healthy training run shows a steep initial drop (as embeddings move from random initialization toward meaningful positions) followed by a gradual leveling off (as the learning rate decreases and the embeddings fine-tune). If the loss plateaus very early and then stays flat, the model may have insufficient capacity for the dataset, or the learning rate may be too low. If the loss oscillates without settling, the initial learning rate may be too high.
Inspecting Learned Embeddings
After training, we extract the center word embeddings and verify they have learned some structure. On a corpus this small, we cannot expect the rich semantic relationships that emerge from training on billions of tokens, but we can check whether frequently co-occurring words have ended up near each other in embedding space.
The cosine similarity function computes the angle between two embedding vectors independent of their magnitude. Words with similar directional orientations in embedding space receive high cosine similarity scores, while words that point in different directions receive low or negative scores. Using normalized vectors for all comparisons ensures that the magnitude of the embedding (which can vary due to how frequently a word appears as a negative sample) does not influence the similarity ranking.
# Extract embeddings as numpy array
embeddings = model_pt.center_embeddings.weight.detach().numpy()
def get_most_similar(query_word, embeddings, word_to_idx, idx_to_word, topn=5):
"""Find most similar words using cosine similarity."""
if query_word not in word_to_idx:
return []
query_vec = embeddings[word_to_idx[query_word]]
query_norm = np.linalg.norm(query_vec) + 1e-10
# Compute cosine similarity with all words
norms = np.linalg.norm(embeddings, axis=1, keepdims=True) + 1e-10
normalized = embeddings / norms
query_normalized = query_vec / query_norm
similarities = normalized @ query_normalized
# Sort and return top results (excluding the query word itself)
ranked = np.argsort(similarities)[::-1]
results = []
for idx in ranked:
word = idx_to_word[idx]
if word != query_word:
results.append((word, float(similarities[idx])))
if len(results) >= topn:
break
return resultsMost similar to 'fox': trains 0.3594 vec 0.3531 equals 0.2815 learn 0.2753 Most similar to 'dog': tokens 0.3206 machine 0.2897 between 0.2707 learning 0.2616 Most similar to 'word': billions 0.4332 vocabularies 0.3515 on 0.2750 man 0.2366 Most similar to 'learn': meaning 0.3962 representations 0.3031 near 0.2769 equals 0.2753
Visualizing the Embedding Space
A 2D t-SNE projection makes it possible to visually verify that the embeddings have learned meaningful structure, with semantically related words clustering together. t-SNE (t-distributed Stochastic Neighbor Embedding) is a nonlinear dimensionality reduction technique that preserves local neighborhoods: words that are close in the high-dimensional embedding space tend to end up close in the 2D projection.
Interpreting t-SNE plots requires some caution. The distances between clusters in the 2D projection do not necessarily reflect the true distances in the original embedding space. t-SNE optimizes for local structure, so two clusters that appear far apart might be moderately similar in the original space. The meaningful signal in a t-SNE plot is which words cluster together, not how far apart different clusters are.
