Skip-gram Model › Implementing Skip-gram from Scratch
Language AI HandbookAll chapters Skip-gram Model › Implementing Skip-gram from ScratchOne section · Sections in this chapter7 / 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
Let us build a complete Skip-gram implementation using PyTorch. We will train on the small corpus above and observe how embeddings develop.
First, set up the training tensors:
import torch
# Convert word pairs to integer indices
center_words = torch.tensor(
[word_to_idx[c] for c, _ in pairs], dtype=torch.long
)
context_words = torch.tensor(
[word_to_idx[ctx] for _, ctx in pairs], dtype=torch.long
)
V = len(vocab) # vocabulary size
D = 15 # embedding dimension
print(f"Vocabulary size V = {V}")
print(f"Embedding dimension D = {D}")
print(f"Number of training pairs: {len(pairs)}")The Skip-gram model architecture has two embedding matrices, one for the center word role and one for the context word role:
import torch.nn as nn
class SkipGram(nn.Module):
def __init__(self, vocab_size, embed_dim):
super().__init__()
# Input embeddings: one row per word, used when word is center
self.center_embeddings = nn.Embedding(vocab_size, embed_dim)
# Context embeddings: one row per word, used when word is context
self.context_embeddings = nn.Embedding(vocab_size, embed_dim)
# Initialize weights: small uniform for center, zeros for context
nn.init.uniform_(
self.center_embeddings.weight, -0.5 / embed_dim, 0.5 / embed_dim
)
nn.init.zeros_(self.context_embeddings.weight)
def get_all_scores(self, center_idx):
"""Compute dot-product scores against the full vocabulary for softmax."""
center_vec = self.center_embeddings(center_idx) # (batch, D)
all_context = self.context_embeddings.weight # (V, D)
return center_vec @ all_context.T # (batch, V)The initialization deserves a brief note. Center embeddings are initialized with small random values so that the model starts with some variation between words. Context embeddings are initialized to zero, which means the initial softmax output is uniform over the vocabulary. This gives a clean starting point before training breaks the symmetry.
Now train the model using full softmax cross-entropy loss. Note that nn.CrossEntropyLoss expects raw logits, not probabilities, so we pass the output of get_all_scores directly without applying softmax explicitly.
import torch.optim as optim
torch.manual_seed(42)
model = SkipGram(V, D)
optimizer = optim.Adam(model.parameters(), lr=0.05)
loss_fn = nn.CrossEntropyLoss()
num_epochs = 500
losses = []
for epoch in range(num_epochs):
optimizer.zero_grad()
logits = model.get_all_scores(center_words) # (N, V)
loss = loss_fn(logits, context_words)
loss.backward()
optimizer.step()
losses.append(loss.item())Training complete over 500 epochs. Initial loss: 3.1781 Final loss: 1.5457 Loss reduction: 1.6323 (51.4%)
The significant loss reduction shows the model has learned to associate center words with their likely context words. The initial loss is close to , which is the entropy of a uniform distribution over the vocabulary, confirming that the model starts with no preferences. By the end of training, the model has concentrated probability mass on the words observed as context, moving well below the random baseline.
Now let us examine the learned embeddings to see whether they reflect meaningful relationships:
# Extract learned center embeddings
with torch.no_grad():
embeddings = model.center_embeddings.weight.numpy()
import numpy as np
# Compute cosine similarity between two words
def cosine_sim(w1, w2, emb, w2i):
v1 = emb[w2i[w1]]
v2 = emb[w2i[w2]]
return float(
np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2) + 1e-8)
)
# Find nearest neighbor for a query word
def nearest_neighbors(query, emb, vocab, w2i, top_k=4):
sims = [(w, cosine_sim(query, w, emb, w2i)) for w in vocab if w != query]
return sorted(sims, key=lambda x: -x[1])[:top_k]
# Compare semantically related and unrelated pairs
pairs_to_check = [
("cat", "dog"), # Both pets appearing in similar contexts
("cat", "mat"), # Co-occur syntactically but different categories
("dog", "park"), # Dog frequently appears near park
("cat", "park"), # Weaker association
]Cosine similarities between word pairs:
---------------------------------------------
cosine('cat', 'dog') = 0.4527
cosine('cat', 'mat') = 0.5794
cosine('dog', 'park') = 0.5781
cosine('cat', 'park') = 0.3836
Nearest neighbors:
'cat': 'mat' (0.579), 'mouse' (0.488), 'friends' (0.473), 'dog' (0.453)
'dog': 'park' (0.578), 'cat' (0.453), 'ran' (0.348), 'at' (0.289)Even on this tiny corpus of six sentences, the model has learned that "cat" and "dog" are more similar to each other than either is to random words. The corpus explicitly connects them: "a dog and a cat can be friends" and "cats and dogs both make good pets" place them in near-identical contexts. On a real corpus with billions of words, these patterns accumulate into fine-grained semantic structure across the full vocabulary.
Key Parameters
The key parameters for the Skip-gram model are:
- embed_dim: Dimensionality of the word vectors. Typical values are 50 to 300. Larger dimensions can capture more nuance but require more data and compute to estimate reliably. If you train on too little data with a high-dimensional embedding, many dimensions will remain random noise. Mikolov et al. used 300 dimensions for their best results on analogy tasks.
- window_size: Controls the context radius. Smaller windows (2 to 3) capture syntactic relationships; larger windows (5 to 10) capture semantic ones. The original Word2Vec paper used 5 for Skip-gram.
- learning_rate: Step size for gradient updates. Adam with learning rates around 0.001 to 0.05 works well. The original implementation used stochastic gradient descent with a learning rate that decays linearly from 0.025 to near zero over training.
- num_epochs / training_steps: More training generally improves quality, with diminishing returns beyond a point determined by corpus size. For large corpora, a single pass through the data is often sufficient.
- min_count: The minimum number of times a word must appear in the corpus to be included in the vocabulary. Rare words have unreliable embeddings because they appear too rarely to accumulate meaningful gradient signal. Typical values are 5 to 10.